wrap_writer_test.go
1package middleware
2
3import (
4 "bytes"
5 "net/http"
6 "net/http/httptest"
7 "testing"
8
9 "github.com/stretchr/testify/assert"
10)
11
12func TestHttpFancyWriterRemembersWroteHeaderWhenFlushed(t *testing.T) {
13 f := &httpFancyWriter{basicWriter: basicWriter{ResponseWriter: httptest.NewRecorder()}}
14 f.Flush()
15
16 if !f.wroteHeader {
17 t.Fatal("want Flush to have set wroteHeader=true")
18 }
19}
20
21func TestHttp2FancyWriterRemembersWroteHeaderWhenFlushed(t *testing.T) {
22 f := &http2FancyWriter{basicWriter{ResponseWriter: httptest.NewRecorder()}}
23 f.Flush()
24
25 if !f.wroteHeader {
26 t.Fatal("want Flush to have set wroteHeader=true")
27 }
28}
29
30func TestBasicWritesTeesWritesWithoutDiscard(t *testing.T) {
31 // explicitly create the struct instead of NewRecorder to control the value of Code
32 original := &httptest.ResponseRecorder{
33 HeaderMap: make(http.Header),
34 Body: new(bytes.Buffer),
35 }
36 wrap := &basicWriter{ResponseWriter: original}
37
38 var buf bytes.Buffer
39 wrap.Tee(&buf)
40
41 _, err := wrap.Write([]byte("hello world"))
42 assert.NoError(t, err)
43
44 assert.Equal(t, 200, original.Code)
45 assert.Equal(t, []byte("hello world"), original.Body.Bytes())
46 assert.Equal(t, []byte("hello world"), buf.Bytes())
47 assert.Equal(t, 11, wrap.BytesWritten())
48}
49
50func TestBasicWriterDiscardsWritesToOriginalResponseWriter(t *testing.T) {
51 t.Run("With Tee", func(t *testing.T) {
52 // explicitly create the struct instead of NewRecorder to control the value of Code
53 original := &httptest.ResponseRecorder{
54 HeaderMap: make(http.Header),
55 Body: new(bytes.Buffer),
56 }
57 wrap := &basicWriter{ResponseWriter: original}
58
59 var buf bytes.Buffer
60 wrap.Tee(&buf)
61 wrap.Discard()
62
63 _, err := wrap.Write([]byte("hello world"))
64 assert.NoError(t, err)
65
66 assert.Equal(t, 0, original.Code) // wrapper shouldn't call WriteHeader implicitly
67 assert.Equal(t, 0, original.Body.Len())
68 assert.Equal(t, []byte("hello world"), buf.Bytes())
69 assert.Equal(t, 11, wrap.BytesWritten())
70 })
71
72 t.Run("Without Tee", func(t *testing.T) {
73 // explicitly create the struct instead of NewRecorder to control the value of Code
74 original := &httptest.ResponseRecorder{
75 HeaderMap: make(http.Header),
76 Body: new(bytes.Buffer),
77 }
78 wrap := &basicWriter{ResponseWriter: original}
79 wrap.Discard()
80
81 _, err := wrap.Write([]byte("hello world"))
82 assert.NoError(t, err)
83
84 assert.Equal(t, 0, original.Code) // wrapper shouldn't call WriteHeader implicitly
85 assert.Equal(t, 0, original.Body.Len())
86 assert.Equal(t, 11, wrap.BytesWritten())
87 })
88}
