1package negroni
2
3import (
4	"bufio"
5	"net"
6	"net/http"
7	"net/http/httptest"
8	"testing"
9	"time"
10)
11
12type closeNotifyingRecorder struct {
13	*httptest.ResponseRecorder
14	closed chan bool
15}
16
17func newCloseNotifyingRecorder() *closeNotifyingRecorder {
18	return &closeNotifyingRecorder{
19		httptest.NewRecorder(),
20		make(chan bool, 1),
21	}
22}
23
24func (c *closeNotifyingRecorder) close() {
25	c.closed <- true
26}
27
28func (c *closeNotifyingRecorder) CloseNotify() <-chan bool {
29	return c.closed
30}
31
32type hijackableResponse struct {
33	Hijacked bool
34}
35
36func newHijackableResponse() *hijackableResponse {
37	return &hijackableResponse{}
38}
39
40func (h *hijackableResponse) Header() http.Header           { return nil }
41func (h *hijackableResponse) Write(buf []byte) (int, error) { return 0, nil }
42func (h *hijackableResponse) WriteHeader(code int)          {}
43func (h *hijackableResponse) Flush()                        {}
44func (h *hijackableResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) {
45	h.Hijacked = true
46	return nil, nil, nil
47}
48
49func TestResponseWriterWritingString(t *testing.T) {
50	rec := httptest.NewRecorder()
51	rw := NewResponseWriter(rec)
52
53	rw.Write([]byte("Hello world"))
54
55	expect(t, rec.Code, rw.Status())
56	expect(t, rec.Body.String(), "Hello world")
57	expect(t, rw.Status(), http.StatusOK)
58	expect(t, rw.Size(), 11)
59	expect(t, rw.Written(), true)
60}
61
62func TestResponseWriterWritingStrings(t *testing.T) {
63	rec := httptest.NewRecorder()
64	rw := NewResponseWriter(rec)
65
66	rw.Write([]byte("Hello world"))
67	rw.Write([]byte("foo bar bat baz"))
68
69	expect(t, rec.Code, rw.Status())
70	expect(t, rec.Body.String(), "Hello worldfoo bar bat baz")
71	expect(t, rw.Status(), http.StatusOK)
72	expect(t, rw.Size(), 26)
73}
74
75func TestResponseWriterWritingHeader(t *testing.T) {
76	rec := httptest.NewRecorder()
77	rw := NewResponseWriter(rec)
78
79	rw.WriteHeader(http.StatusNotFound)
80
81	expect(t, rec.Code, rw.Status())
82	expect(t, rec.Body.String(), "")
83	expect(t, rw.Status(), http.StatusNotFound)
84	expect(t, rw.Size(), 0)
85}
86
87func TestResponseWriterBefore(t *testing.T) {
88	rec := httptest.NewRecorder()
89	rw := NewResponseWriter(rec)
90	result := ""
91
92	rw.Before(func(ResponseWriter) {
93		result += "foo"
94	})
95	rw.Before(func(ResponseWriter) {
96		result += "bar"
97	})
98
99	rw.WriteHeader(http.StatusNotFound)
100
101	expect(t, rec.Code, rw.Status())
102	expect(t, rec.Body.String(), "")
103	expect(t, rw.Status(), http.StatusNotFound)
104	expect(t, rw.Size(), 0)
105	expect(t, result, "barfoo")
106}
107
108func TestResponseWriterHijack(t *testing.T) {
109	hijackable := newHijackableResponse()
110	rw := NewResponseWriter(hijackable)
111	hijacker, ok := rw.(http.Hijacker)
112	expect(t, ok, true)
113	_, _, err := hijacker.Hijack()
114	if err != nil {
115		t.Error(err)
116	}
117	expect(t, hijackable.Hijacked, true)
118}
119
120func TestResponseWriteHijackNotOK(t *testing.T) {
121	hijackable := new(http.ResponseWriter)
122	rw := NewResponseWriter(*hijackable)
123	hijacker, ok := rw.(http.Hijacker)
124	expect(t, ok, true)
125	_, _, err := hijacker.Hijack()
126
127	refute(t, err, nil)
128}
129
130func TestResponseWriterCloseNotify(t *testing.T) {
131	rec := newCloseNotifyingRecorder()
132	rw := NewResponseWriter(rec)
133	closed := false
134	notifier := rw.(http.CloseNotifier).CloseNotify()
135	rec.close()
136	select {
137	case <-notifier:
138		closed = true
139	case <-time.After(time.Second):
140	}
141	expect(t, closed, true)
142}
143
144func TestResponseWriterFlusher(t *testing.T) {
145	rec := httptest.NewRecorder()
146	rw := NewResponseWriter(rec)
147
148	_, ok := rw.(http.Flusher)
149	expect(t, ok, true)
150}
151