From 0406abe066d756b85766e3d7eb110a2cbe157c67 Mon Sep 17 00:00:00 2001
From: Rashad Ansari <rashad.ansari1996@gmail.com>
Date: Wed, 2 Dec 2020 15:43:42 +0330
Subject: [PATCH] Add the ability to change the status code using Response
 beforeFuncs

---
 response.go      |  4 ++--
 response_test.go | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/response.go b/response.go
index ca7405c5d..84f7c9e7e 100644
--- a/response.go
+++ b/response.go
@@ -56,11 +56,11 @@ func (r *Response) WriteHeader(code int) {
 		r.echo.Logger.Warn("response already committed")
 		return
 	}
+	r.Status = code
 	for _, fn := range r.beforeFuncs {
 		fn()
 	}
-	r.Status = code
-	r.Writer.WriteHeader(code)
+	r.Writer.WriteHeader(r.Status)
 	r.Committed = true
 }
 
diff --git a/response_test.go b/response_test.go
index 7a9c51c66..d95e079f9 100644
--- a/response_test.go
+++ b/response_test.go
@@ -56,3 +56,19 @@ func TestResponse_Flush(t *testing.T) {
 	res.Flush()
 	assert.True(t, rec.Flushed)
 }
+
+func TestResponse_ChangeStatusCodeBeforeWrite(t *testing.T) {
+	e := New()
+	rec := httptest.NewRecorder()
+	res := &Response{echo: e, Writer: rec}
+
+	res.Before(func() {
+		if 200 < res.Status && res.Status < 300 {
+			res.Status = 200
+		}
+	})
+
+	res.WriteHeader(209)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+}