|
9 | 9 | "net/http" |
10 | 10 | "net/http/httptest" |
11 | 11 | "strconv" |
| 12 | + "strings" |
12 | 13 | "testing" |
13 | 14 |
|
14 | 15 | "github.com/stretchr/testify/assert" |
@@ -173,30 +174,75 @@ func TestBuffer_requestLimitReached(t *testing.T) { |
173 | 174 | } |
174 | 175 |
|
175 | 176 | func TestBuffer_responseLimitReached(t *testing.T) { |
176 | | - srv := testutils.NewHandler(func(w http.ResponseWriter, _ *http.Request) { |
177 | | - _, _ = w.Write([]byte("hello, this response is too large")) |
178 | | - }) |
179 | | - t.Cleanup(srv.Close) |
180 | | - |
181 | | - // forwarder will proxy the request to whatever destination |
182 | | - fwd := forward.New(false) |
183 | | - |
184 | | - // this is our redirect to server |
185 | | - rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { |
186 | | - req.URL = testutils.MustParseRequestURI(srv.URL) |
187 | | - fwd.ServeHTTP(w, req) |
188 | | - }) |
189 | | - |
190 | | - // stream handler will forward requests to redirect |
191 | | - st, err := New(rdr, MaxResponseBodyBytes(4)) |
192 | | - require.NoError(t, err) |
193 | | - |
194 | | - proxy := httptest.NewServer(st) |
195 | | - t.Cleanup(proxy.Close) |
196 | | - |
197 | | - re, _, err := testutils.Get(proxy.URL) |
198 | | - require.NoError(t, err) |
199 | | - assert.Equal(t, http.StatusInternalServerError, re.StatusCode) |
| 177 | + cases := []struct { |
| 178 | + name string |
| 179 | + body string |
| 180 | + maxResponseBodyBytes int64 |
| 181 | + }{ |
| 182 | + { |
| 183 | + name: "small limit with body larger than max response bytes", |
| 184 | + body: "hello, this response is too large", |
| 185 | + maxResponseBodyBytes: 4, |
| 186 | + }, |
| 187 | + { |
| 188 | + name: "small limit with body larger than 32768 bytes", |
| 189 | + body: strings.Repeat("A", 32769), |
| 190 | + maxResponseBodyBytes: 4, |
| 191 | + }, |
| 192 | + { |
| 193 | + name: "larger limit with body larger than 32768 bytes", |
| 194 | + body: strings.Repeat("A", 32769), |
| 195 | + maxResponseBodyBytes: 2000, |
| 196 | + }, |
| 197 | + { |
| 198 | + name: "larger limit with body larger than 32768 + 1999 bytes", |
| 199 | + body: strings.Repeat("A", 32769+1999), |
| 200 | + maxResponseBodyBytes: 2000, |
| 201 | + }, |
| 202 | + { |
| 203 | + name: "larger limit with body larger than 32768 + 2000 bytes", |
| 204 | + body: strings.Repeat("A", 32769+2000), |
| 205 | + maxResponseBodyBytes: 2000, |
| 206 | + }, |
| 207 | + { |
| 208 | + name: "larger limit with body larger than 65536 + 1999 bytes", |
| 209 | + body: strings.Repeat("A", 65537+1999), |
| 210 | + maxResponseBodyBytes: 2000, |
| 211 | + }, |
| 212 | + { |
| 213 | + name: "larger limit with body larger than 65536 + 2000 bytes", |
| 214 | + body: strings.Repeat("A", 65537+2000), |
| 215 | + maxResponseBodyBytes: 2000, |
| 216 | + }, |
| 217 | + } |
| 218 | + for _, tc := range cases { |
| 219 | + t.Run(tc.name, func(t *testing.T) { |
| 220 | + srv := testutils.NewHandler(func(w http.ResponseWriter, _ *http.Request) { |
| 221 | + _, _ = w.Write([]byte(tc.body)) |
| 222 | + }) |
| 223 | + t.Cleanup(srv.Close) |
| 224 | + |
| 225 | + // forwarder will proxy the request to whatever destination |
| 226 | + fwd := forward.New(false) |
| 227 | + |
| 228 | + // this is our redirect to server |
| 229 | + rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { |
| 230 | + req.URL = testutils.MustParseRequestURI(srv.URL) |
| 231 | + fwd.ServeHTTP(w, req) |
| 232 | + }) |
| 233 | + |
| 234 | + // stream handler will forward requests to redirect |
| 235 | + st, err := New(rdr, MaxResponseBodyBytes(tc.maxResponseBodyBytes)) |
| 236 | + require.NoError(t, err) |
| 237 | + |
| 238 | + proxy := httptest.NewServer(st) |
| 239 | + t.Cleanup(proxy.Close) |
| 240 | + |
| 241 | + re, _, err := testutils.Get(proxy.URL) |
| 242 | + require.NoError(t, err) |
| 243 | + assert.Equal(t, http.StatusInternalServerError, re.StatusCode) |
| 244 | + }) |
| 245 | + } |
200 | 246 | } |
201 | 247 |
|
202 | 248 | func TestBuffer_fileStreamingResponse(t *testing.T) { |
|
0 commit comments