chore: avoid nested WithCodeResponseWriter (#3406)

rest/handler/timeouthandler_test.go

@@ -100,6 +100,18 @@ func TestWithinTimeout(t *testing.T) {
 	assert.Equal(t, http.StatusOK, resp.Code)
 }
 
+func TestWithinTimeoutBadCode(t *testing.T) {
+	timeoutHandler := TimeoutHandler(time.Second)
+	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, http.StatusInternalServerError, resp.Code)
+}
+
 func TestWithTimeoutTimedout(t *testing.T) {
 	timeoutHandler := TimeoutHandler(time.Millisecond)
 	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -208,9 +220,7 @@ func TestTimeoutHijack(t *testing.T) {
 	resp := httptest.NewRecorder()
 
 	writer := &timeoutWriter{
-		w: &response.WithCodeResponseWriter{
-			Writer: resp,
-		},
+		w: response.NewWithCodeResponseWriter(resp),
 	}
 
 	assert.NotPanics(t, func() {
@@ -218,9 +228,7 @@ func TestTimeoutHijack(t *testing.T) {
 	})
 
 	writer = &timeoutWriter{
-		w: &response.WithCodeResponseWriter{
-			Writer: mockedHijackable{resp},
-		},
+		w: response.NewWithCodeResponseWriter(mockedHijackable{resp}),
 	}
 
 	assert.NotPanics(t, func() {
@@ -274,9 +282,7 @@ func TestTimeoutWriter_Hijack(t *testing.T) {
 func TestTimeoutWroteTwice(t *testing.T) {
 	c := logtest.NewCollector(t)
 	writer := &timeoutWriter{
-		w: &response.WithCodeResponseWriter{
-			Writer: httptest.NewRecorder(),
-		},
+		w:   response.NewWithCodeResponseWriter(httptest.NewRecorder()),
 		h:   make(http.Header),
 		req: httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody),
 	}