|
| 1 | +// +build go1.7 |
| 2 | + |
| 3 | +package gin |
| 4 | + |
| 5 | +import ( |
| 6 | + "bytes" |
| 7 | + "net/http" |
| 8 | + "net/http/httptest" |
| 9 | + "testing" |
| 10 | + |
| 11 | + "github.com/stretchr/testify/assert" |
| 12 | +) |
| 13 | + |
| 14 | +type interceptedWriter struct { |
| 15 | + ResponseWriter |
| 16 | + b *bytes.Buffer |
| 17 | +} |
| 18 | + |
| 19 | +func (i interceptedWriter) WriteHeader(code int) { |
| 20 | + i.Header().Del("X-Test") |
| 21 | + i.ResponseWriter.WriteHeader(code) |
| 22 | +} |
| 23 | +func TestInterceptedHeader(t *testing.T) { |
| 24 | + w := httptest.NewRecorder() |
| 25 | + c, r := CreateTestContext(w) |
| 26 | + |
| 27 | + r.Use(func(c *Context) { |
| 28 | + i := interceptedWriter{ |
| 29 | + ResponseWriter: c.Writer, |
| 30 | + b: bytes.NewBuffer(nil), |
| 31 | + } |
| 32 | + c.Writer = i |
| 33 | + c.Next() |
| 34 | + c.Header("X-Test", "overridden") |
| 35 | + c.Writer = i.ResponseWriter |
| 36 | + }) |
| 37 | + r.GET("/", func(c *Context) { |
| 38 | + c.Header("X-Test", "original") |
| 39 | + c.Header("X-Test-2", "present") |
| 40 | + c.String(http.StatusOK, "hello world") |
| 41 | + }) |
| 42 | + c.Request = httptest.NewRequest("GET", "/", nil) |
| 43 | + r.HandleContext(c) |
| 44 | + // Result() has headers frozen when WriteHeaderNow() has been called |
| 45 | + // Compared to this time, this is when the response headers will be flushed |
| 46 | + // As response is flushed on c.String, the Header cannot be set by the first |
| 47 | + // middleware. Assert this |
| 48 | + assert.Equal(t, "", w.Result().Header.Get("X-Test")) |
| 49 | + assert.Equal(t, "present", w.Result().Header.Get("X-Test-2")) |
| 50 | +} |
0 commit comments