fix: handle missing x-response-json

main.go

@@ -32,42 +32,54 @@ func main() {
 }
 
 func handler(w http.ResponseWriter, r *http.Request) {
-	var (
-		contentType  = "application/json; charset=utf-8"
-		responseBody = r.Header["X-Response-Json"][0]
-		responseCode int
-	)
+	responseBody, responseCode := buildResponse(r.Header)
 
-	responseCode, err := validateResponseCode(r.Header)
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.WriteHeader(responseCode)
+	w.Write([]byte(responseBody))
+}
+
+func buildResponse(header map[string][]string) (string, int) {
+	responseBody, err := validateResponseBody(header)
 	if err != nil {
-		responseBody = err.Error()
-	} else if !isJSON(responseBody) {
-		responseBody = errorResponseFormatter("x-response-json must be valid JSON").Error()
-		responseCode = http.StatusBadRequest
+		return err.Error(), http.StatusBadRequest
+
 	}
 
-	w.Header().Set("Content-Type", contentType)
-	w.WriteHeader(responseCode)
-	w.Write([]byte(responseBody))
+	responseCode, err := validateResponseCode(header)
+	if err != nil {
+		return err.Error(), http.StatusBadRequest
+	}
+
+	return responseBody, responseCode
+}
+
+func validateResponseBody(header map[string][]string) (string, error) {
+	if _, exists := header["X-Response-Json"]; !exists {
+		return "", errorResponseFormatter("x-response-json must be set on the request")
+	}
+
+	if !isJSON(header["X-Response-Json"][0]) {
+		return "", errorResponseFormatter("x-response-json must be valid JSON")
+	}
+
+	return header["X-Response-Json"][0], nil
 }
 
 func validateResponseCode(header map[string][]string) (int, error) {
 	if val, exists := header["X-Response-Code"]; exists {
-		// verify its a number
 		code, err := strconv.Atoi(val[0])
 		if err != nil {
 			return http.StatusBadRequest, errorResponseFormatter("x-response-code must be a number")
 		}
 
-		// verify it falls in the range of status codes
 		if !(code >= 100 && code <= 599) {
 			return http.StatusBadRequest, errorResponseFormatter("x-response-code must be between 100 and 599")
 		}
 
 		return code, nil
 	}
 
-	// if x-response-code was not supplied default to a 200
 	return http.StatusOK, nil
 }
 

main_test.go

@@ -41,6 +41,32 @@ func TestHandlerValidRequest(t *testing.T) {
 	}
 }
 
+func TestHandlerWithoutXResponseJson(t *testing.T) {
+	expectedResponseBody := `{"error":"x-response-json must be set on the request"}`
+	expectedResponseCode := http.StatusBadRequest
+
+	// set up the request
+	req, err := http.NewRequest("GET", "/test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// make the request
+	rr := httptest.NewRecorder()
+	h := http.HandlerFunc(handler)
+	h.ServeHTTP(rr, req)
+
+	// status code
+	if status := rr.Code; status != expectedResponseCode {
+		t.Errorf("handler returned wrong status code: got '%d' want '%d'", status, expectedResponseCode)
+	}
+
+	// response body
+	if rr.Body.String() != expectedResponseBody {
+		t.Errorf("handler returned unexpected body: got '%s' want '%s'", rr.Body.String(), expectedResponseBody)
+	}
+}
+
 func TestHandlerWithInvalidXResponseJson(t *testing.T) {
 	invalidResponseBody := `{"foo":bar}`
 	expectedResponseBody := `{"error":"x-response-json must be valid JSON"}`