diff --git a/go-api/main.go b/go-api/main.go
index f309b33..f560e30 100644
--- a/go-api/main.go
+++ b/go-api/main.go
@@ -10,22 +10,24 @@ import (
 func handler(w http.ResponseWriter, r *http.Request) {
 	fmt.Println("Running handler...")
 
-	// If the `force-fail` header is set, get its value and use it as the status
-	// code.
+	w.WriteHeader(getResponseCode(r))
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(map[string]any{})
+}
+
+func getResponseCode(r *http.Request) int {
+	// If the `force-fail` header is set, get and return its value.
 	if failCode := r.Header.Get("force-fail"); failCode != "" {
 		fmt.Println("`force-fail` header set...")
 
 		if code, err := strconv.Atoi(failCode); err == nil {
 			fmt.Printf("Setting the response code to %d...", code)
 
-			w.WriteHeader(code)
+			return code
 		}
-	} else {
-		w.WriteHeader(http.StatusOK)
 	}
 
-	w.Header().Set("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(map[string]any{})
+	return http.StatusOK
 }
 
 func main() {