@@ -18,7 +18,6 @@ import (
1818 "github.com/coder/coder/coderd/httpmw"
1919 "github.com/coder/coder/coderd/rbac"
2020 "github.com/coder/coder/codersdk"
21- "github.com/coder/coder/testutil"
2221)
2322
2423func randRemoteAddr () string {
@@ -34,38 +33,39 @@ func TestRateLimit(t *testing.T) {
3433 t .Run ("NoUserSucceeds" , func (t * testing.T ) {
3534 t .Parallel ()
3635 rtr := chi .NewRouter ()
37- rtr .Use (httpmw .RateLimit (5 , time .Second ))
36+ rtr .Use (httpmw .RateLimit (1 , time .Second ))
3837 rtr .Get ("/" , func (rw http.ResponseWriter , r * http.Request ) {
3938 rw .WriteHeader (http .StatusOK )
4039 })
4140
42- require . Eventually ( t , func () bool {
41+ for i := 0 ; i < 5 ; i ++ {
4342 req := httptest .NewRequest ("GET" , "/" , nil )
4443 rec := httptest .NewRecorder ()
4544 rtr .ServeHTTP (rec , req )
4645 resp := rec .Result ()
47- defer resp .Body .Close ()
48- return resp .StatusCode == http .StatusTooManyRequests
49- }, testutil . WaitShort , testutil . IntervalFast )
46+ _ = resp .Body .Close ()
47+ require . Equal ( t , i != 0 , resp .StatusCode == http .StatusTooManyRequests )
48+ }
5049 })
5150
5251 t .Run ("RandomIPs" , func (t * testing.T ) {
5352 t .Parallel ()
5453 rtr := chi .NewRouter ()
55- rtr .Use (httpmw .RateLimit (5 , time .Second ))
54+ // Because these are random IPs, the limit should never be hit!
55+ rtr .Use (httpmw .RateLimit (1 , time .Second ))
5656 rtr .Get ("/" , func (rw http.ResponseWriter , r * http.Request ) {
5757 rw .WriteHeader (http .StatusOK )
5858 })
5959
60- require . Never ( t , func () bool {
60+ for i := 0 ; i < 5 ; i ++ {
6161 req := httptest .NewRequest ("GET" , "/" , nil )
6262 rec := httptest .NewRecorder ()
6363 req .RemoteAddr = randRemoteAddr ()
6464 rtr .ServeHTTP (rec , req )
6565 resp := rec .Result ()
66- defer resp .Body .Close ()
67- return resp .StatusCode == http .StatusTooManyRequests
68- }, testutil . WaitShort , testutil . IntervalFast )
66+ _ = resp .Body .Close ()
67+ require . False ( t , resp .StatusCode == http .StatusTooManyRequests )
68+ }
6969 })
7070
7171 t .Run ("RegularUser" , func (t * testing.T ) {
@@ -81,7 +81,7 @@ func TestRateLimit(t *testing.T) {
8181 Optional : false ,
8282 }))
8383
84- rtr .Use (httpmw .RateLimit (5 , time .Second ))
84+ rtr .Use (httpmw .RateLimit (1 , time .Second ))
8585 rtr .Get ("/" , func (rw http.ResponseWriter , r * http.Request ) {
8686 rw .WriteHeader (http .StatusOK )
8787 })
@@ -98,17 +98,17 @@ func TestRateLimit(t *testing.T) {
9898 defer resp .Body .Close ()
9999 require .Equal (t , http .StatusPreconditionRequired , resp .StatusCode )
100100
101- require . Eventually ( t , func () bool {
101+ for i := 0 ; i < 5 ; i ++ {
102102 req := httptest .NewRequest ("GET" , "/" , nil )
103103 req .Header .Set (codersdk .SessionTokenHeader , key )
104104 rec := httptest .NewRecorder ()
105105 // Assert we're not using IP address.
106106 req .RemoteAddr = randRemoteAddr ()
107107 rtr .ServeHTTP (rec , req )
108108 resp := rec .Result ()
109- defer resp .Body .Close ()
110- return resp .StatusCode == http .StatusTooManyRequests
111- }, testutil . WaitShort , testutil . IntervalFast )
109+ _ = resp .Body .Close ()
110+ require . Equal ( t , i != 0 , resp .StatusCode == http .StatusTooManyRequests )
111+ }
112112 })
113113
114114 t .Run ("OwnerBypass" , func (t * testing.T ) {
@@ -127,12 +127,12 @@ func TestRateLimit(t *testing.T) {
127127 Optional : false ,
128128 }))
129129
130- rtr .Use (httpmw .RateLimit (5 , time .Second ))
130+ rtr .Use (httpmw .RateLimit (1 , time .Second ))
131131 rtr .Get ("/" , func (rw http.ResponseWriter , r * http.Request ) {
132132 rw .WriteHeader (http .StatusOK )
133133 })
134134
135- require . Never ( t , func () bool {
135+ for i := 0 ; i < 5 ; i ++ {
136136 req := httptest .NewRequest ("GET" , "/" , nil )
137137 req .Header .Set (codersdk .SessionTokenHeader , key )
138138 req .Header .Set (codersdk .BypassRatelimitHeader , "true" )
@@ -141,8 +141,8 @@ func TestRateLimit(t *testing.T) {
141141 req .RemoteAddr = randRemoteAddr ()
142142 rtr .ServeHTTP (rec , req )
143143 resp := rec .Result ()
144- defer resp .Body .Close ()
145- return resp .StatusCode == http .StatusTooManyRequests
146- }, testutil . WaitShort , testutil . IntervalFast )
144+ _ = resp .Body .Close ()
145+ require . False ( t , resp .StatusCode == http .StatusTooManyRequests )
146+ }
147147 })
148148}
0 commit comments