Handling all no content status codes in a generic way

runtime/constants.go

@@ -137,3 +137,8 @@ var knownStatusCodes = map[int]bool{
 	http.StatusNotExtended:                   true, // 510
 	http.StatusNetworkAuthenticationRequired: true, // 511
 }
+
+var noContentStatusCodes = map[int]bool{
+	http.StatusNoContent:   true, //204
+	http.StatusNotModified: true, //304
+}

runtime/server_http_response.go

@@ -292,7 +292,7 @@ func (res *ServerHTTPResponse) flush(ctx context.Context) {
 
 	res.flushed = true
 	res.writeHeader(res.pendingStatusCode)
-	if res.pendingStatusCode != http.StatusNoContent {
+	if _, noContent := noContentStatusCodes[res.pendingStatusCode]; !noContent {
 		res.writeBytes(res.pendingBodyBytes)
 	}
 	res.finish(ctx)

runtime/server_http_response_test.go

@@ -25,6 +25,7 @@ import (
 	"encoding/json"
 	"io/ioutil"
 	"net/http"
+	"strconv"
 	"testing"
 
 	"github.com/buger/jsonparser"
@@ -583,6 +584,81 @@ func TestPendingResponseBody204StatusNoContent(t *testing.T) {
 		"",
 		string(bytes),
 	)
+	assert.Equal(
+		t,
+		"0",
+		strconv.Itoa(len(bytes)),
+	)
+}
+
+func TestPendingResponseBody304StatusNoContent(t *testing.T) {
+	gateway, err := benchGateway.CreateGateway(
+		defaultTestConfig,
+		defaultTestOptions,
+		exampleGateway.CreateGateway,
+	)
+
+	if !assert.NoError(t, err) {
+		return
+	}
+	defer gateway.Close()
+
+	bgateway := gateway.(*benchGateway.BenchGateway)
+	deps := createDefaultDependencies(bgateway)
+	err = bgateway.ActualGateway.HTTPRouter.Handle(
+		"GET", "/foo", http.HandlerFunc(zanzibar.NewRouterEndpoint(
+			bgateway.ActualGateway.ContextExtractor,
+			deps,
+			"foo", "foo",
+			func(
+				ctx context.Context,
+				req *zanzibar.ServerHTTPRequest,
+				res *zanzibar.ServerHTTPResponse,
+			) {
+				obj := &MyBody{
+					Token: "myToken",
+					Client: MyBodyClient{
+						Token: "myClientToken",
+					},
+				}
+				bytes, err := json.Marshal(obj)
+				statusCode := 304
+				assert.NoError(t, err)
+				res.WriteJSON(statusCode, nil, obj)
+
+				pendingBytes, pendingStatusCode := res.GetPendingResponse()
+				assert.Equal(t, bytes, pendingBytes)
+				assert.Equal(t, statusCode, pendingStatusCode)
+
+				headers := res.Headers()
+				assert.NotNil(t, headers)
+			},
+		).HandleRequest),
+	)
+	assert.NoError(t, err)
+
+	resp, err := gateway.MakeRequest("GET", "/foo", nil, nil)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	assert.Equal(t, resp.StatusCode, 304)
+	bytes, err := ioutil.ReadAll(resp.Body)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	//The body would become blank
+	assert.Equal(
+		t,
+		``,
+		string(bytes),
+	)
+	assert.Equal(
+		t,
+		"0",
+		strconv.Itoa(len(bytes)),
+	)
 }
 
 func TestPendingResponseObject(t *testing.T) {