From ac6db4f3bc86c8dd92a8cb09e8a0cc3e037740bb Mon Sep 17 00:00:00 2001
From: Nate Meyer <672246+notnmeyer@users.noreply.github.com>
Date: Mon, 23 Oct 2023 14:32:03 -0700
Subject: [PATCH] handle missing x-response-json

handle missing x-response-json
---
 main.go      | 38 ++++++++++++++++++++++++++++++++++----
 main_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index 50b1c11..f6e3acf 100644
--- a/main.go
+++ b/main.go
@@ -34,23 +34,53 @@ func main() {
 func handler(w http.ResponseWriter, r *http.Request) {
 	var (
 		contentType  = "application/json; charset=utf-8"
-		responseBody = r.Header["X-Response-Json"][0]
+		responseBody string
 		responseCode int
 	)
 
-	responseCode, err := validateResponseCode(r.Header)
+	err := validateResponseBody(r.Header)
+	if err == nil {
+		responseBody = r.Header["X-Response-Json"][0]
+	} else {
+		responseBody = err.Error()
+		responseCode = http.StatusBadRequest
+		writeRequest(w, contentType, responseBody, responseCode)
+		return
+	}
+
+	responseCode, err = validateResponseCode(r.Header)
 	if err != nil {
 		responseBody = err.Error()
-	} else if !isJSON(responseBody) {
-		responseBody = errorResponseFormatter("x-response-json must be valid JSON").Error()
 		responseCode = http.StatusBadRequest
+		writeRequest(w, contentType, responseBody, responseCode)
+		return
 	}
 
+	writeRequest(w, contentType, responseBody, responseCode)
+
+	// w.Header().Set("Content-Type", contentType)
+	// w.WriteHeader(responseCode)
+	// w.Write([]byte(responseBody))
+}
+
+func writeRequest(w http.ResponseWriter, contentType string, responseBody string, responseCode int) {
 	w.Header().Set("Content-Type", contentType)
 	w.WriteHeader(responseCode)
 	w.Write([]byte(responseBody))
 }
 
+func validateResponseBody(header map[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 nil
+}
+
 func validateResponseCode(header map[string][]string) (int, error) {
 	if val, exists := header["X-Response-Code"]; exists {
 		// verify its a number
diff --git a/main_test.go b/main_test.go
index a67510f..236422d 100644
--- a/main_test.go
+++ b/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"}`