Address review comments

Signed-off-by: albertteoh <albert.teoh@logz.io>

cmd/query/app/http_handler.go

@@ -460,9 +460,7 @@ func (aH *APIHandler) writeJSON(w http.ResponseWriter, r *http.Request, response
 		return
 	}
 	w.Header().Set("Content-Type", "application/json")
-	if b, err := w.Write(resp); err != nil {
+	if _, err := w.Write(resp); err != nil {
 		aH.handleError(w, fmt.Errorf("failed writing HTTP response: %w", err), http.StatusInternalServerError)
-	} else {
-		aH.logger.Debug("Successfully wrote HTTP response", zap.Int("bytes", b))
 	}
 }

cmd/query/app/http_handler_test.go

@@ -186,8 +186,7 @@ func TestLogOnServerError(t *testing.T) {
 	assert.Equal(t, e, (*l.logs)[0].f[0].Interface)
 }
 
-// httpResponseErrWriter implements the http.ResponseWriter interface
-// that returns an error on Write.
+// httpResponseErrWriter implements the http.ResponseWriter interface that returns an error on Write.
 type httpResponseErrWriter struct{}
 
 func (h *httpResponseErrWriter) Write([]byte) (int, error) {
@@ -197,13 +196,6 @@ func (h *httpResponseErrWriter) WriteHeader(statusCode int) {}
 func (h *httpResponseErrWriter) Header() http.Header {
 	return http.Header{}
 }
-
-// ErrHandlerFunc injects the httpResponseErrWriter into the HTTP handler.
-type ErrHandlerFunc func(http.ResponseWriter, *http.Request)
-
-func (f ErrHandlerFunc) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	f(&httpResponseErrWriter{}, r)
-}
 func TestWriteJSON(t *testing.T) {
 	testCases := []struct {
 		name         string
@@ -250,19 +242,17 @@ func TestWriteJSON(t *testing.T) {
 	}
 
 	for _, testCase := range testCases {
-		t.Run(testCase.param, func(t *testing.T) {
+		t.Run(testCase.name, func(t *testing.T) {
 			apiHandler := &APIHandler{
 				logger: zap.NewNop(),
 			}
 			httpHandler := func(w http.ResponseWriter, r *http.Request) {
+				if testCase.httpWriteErr {
+					w = &httpResponseErrWriter{}
+				}
 				apiHandler.writeJSON(w, r, &testCase.data)
 			}
-			var server *httptest.Server
-			if testCase.httpWriteErr {
-				server = httptest.NewServer(ErrHandlerFunc(httpHandler))
-			} else {
-				server = httptest.NewServer(http.HandlerFunc(httpHandler))
-			}
+			server := httptest.NewServer(http.HandlerFunc(httpHandler))
 			defer server.Close()
 
 			out := get(server.URL + testCase.param)