From cbb96f4955c3649d5e2ba6e73273a1d982f26f86 Mon Sep 17 00:00:00 2001
From: EwenQuim <ewen.quimerch@gmail.com>
Date: Wed, 18 Dec 2024 08:46:53 +0100
Subject: [PATCH] Set status code before returning when no data

---
 serve.go      |  8 ++++----
 serve_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/serve.go b/serve.go
index 4354ddfe..c6a03979 100644
--- a/serve.go
+++ b/serve.go
@@ -139,14 +139,14 @@ func HTTPHandler[ReturnType, Body any, Contextable ctx[Body]](s *Server, control
 		}
 		w.Header().Add("Server-Timing", Timing{"controller", time.Since(timeController), ""}.String())
 
-		if reflect.TypeOf(ans) == nil {
-			return
-		}
-
 		if route.DefaultStatusCode != 0 {
 			w.WriteHeader(route.DefaultStatusCode)
 		}
 
+		if reflect.TypeOf(ans) == nil {
+			return
+		}
+
 		// TRANSFORM OUT
 		timeTransformOut := time.Now()
 		ans, err = transformOut(r.Context(), ans)
diff --git a/serve_test.go b/serve_test.go
index 209ab156..fdd41498 100644
--- a/serve_test.go
+++ b/serve_test.go
@@ -607,3 +607,39 @@ func newTLSTestHelper() (*tlsTestHelper, error) {
 	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: privateKeyBytes})
 	return &tlsTestHelper{cert: certPEM, key: keyPEM}, nil
 }
+
+func TestReturn204WhenNoData(t *testing.T) {
+	t.Run("can return 204 when no data is being sent", func(t *testing.T) {
+		s := NewServer()
+
+		Get(s, "/", func(_ ContextNoBody) (any, error) {
+			return nil, nil
+		},
+			OptionDefaultStatusCode(204),
+		)
+
+		r := httptest.NewRequest("GET", "/", nil)
+		w := httptest.NewRecorder()
+
+		s.Mux.ServeHTTP(w, r)
+
+		require.Equal(t, 204, w.Code)
+	})
+
+	t.Run("must return 500 when an error is being sent, even with no body", func(t *testing.T) {
+		s := NewServer()
+
+		Get(s, "/", func(_ ContextNoBody) (any, error) {
+			return nil, errors.New("error")
+		},
+			OptionDefaultStatusCode(204),
+		)
+
+		r := httptest.NewRequest("GET", "/", nil)
+		w := httptest.NewRecorder()
+
+		s.Mux.ServeHTTP(w, r)
+
+		require.Equal(t, 500, w.Code)
+	})
+}