feat: add Status method to Writer struct and update test (#63)

- Add a `Status` method to the `Writer` struct
- Modify the `TestWriter_Status` test function to include the new `Status` method

fixed by @zhyee

fixed #52
fixed #51

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

writer.go

@@ -72,6 +72,16 @@ func (w *Writer) FreeBuffer() {
  w.body = nil
 }
 
+// Status we must override Status func here,
+// or the http status code returned by gin.Context.Writer.Status()
+// will always be 200 in other custom gin middlewares.
+func (w *Writer) Status() int {
+ if w.code == 0 || w.timeout {
+ return w.ResponseWriter.Status()
+ }
+ return w.code
+}
+
 func checkWriteHeaderCode(code int) {
  if code < 100 || code > 999 {
  panic(fmt.Sprintf("invalid http status code: %d", code))

writer_test.go

@@ -2,8 +2,13 @@ package timeout
 
 import (
  "fmt"
+ "net/http"
+ "net/http/httptest"
+ "strconv"
  "testing"
+ "time"
 
+ "github.com/gin-gonic/gin"
  "github.com/stretchr/testify/assert"
 )
 
@@ -21,3 +26,34 @@ func TestWriteHeader(t *testing.T) {
  writer.WriteHeader(code2)
  })
 }
+
+func TestWriter_Status(t *testing.T) {
+ r := gin.New()
+
+ r.Use(New(
+ WithTimeout(1*time.Second),
+ WithHandler(func(c *gin.Context) {
+ c.Next()
+ }),
+ WithResponse(testResponse),
+ ))
+
+ r.Use(func(c *gin.Context) {
+ c.Next()
+ statusInMW := c.Writer.Status()
+ c.Request.Header.Set("X-Status-Code-MW-Set", strconv.Itoa(statusInMW))
+ t.Logf("[%s] %s %s %d\n", time.Now().Format(time.RFC3339), c.Request.Method, c.Request.URL, statusInMW)
+ })
+
+ r.GET("/test", func(c *gin.Context) {
+ c.Writer.WriteHeader(http.StatusInternalServerError)
+ })
+
+ w := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodGet, "/test", nil)
+
+ r.ServeHTTP(w, req)
+
+ assert.Equal(t, http.StatusInternalServerError, w.Code)
+ assert.Equal(t, strconv.Itoa(http.StatusInternalServerError), req.Header.Get("X-Status-Code-MW-Set"))
+}