chore: improve times of ratelimit tests (coder#6346)

kylecarbs · web-flow · commit 17adfd1134bb · 2023-02-25T22:01:01.000Z
From 5s to 130ms!
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -484,6 +484,9 @@ func createAnotherUserRetry(t *testing.T, client *codersdk.Client, organizationI
 
 	other := codersdk.New(client.URL)
 	other.SetSessionToken(login.SessionToken)
+	t.Cleanup(func() {
+		other.HTTPClient.CloseIdleConnections()
+	})
 
 	if len(roles) > 0 {
 		// Find the roles for the org vs the site wide roles
diff --git a/coderd/httpmw/ratelimit_test.go b/coderd/httpmw/ratelimit_test.go
@@ -18,7 +18,6 @@ import (
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/rbac"
 	"github.com/coder/coder/codersdk"
-	"github.com/coder/coder/testutil"
 )
 
 func randRemoteAddr() string {
@@ -34,38 +33,39 @@ func TestRateLimit(t *testing.T) {
 	t.Run("NoUserSucceeds", func(t *testing.T) {
 		t.Parallel()
 		rtr := chi.NewRouter()
-		rtr.Use(httpmw.RateLimit(5, time.Second))
+		rtr.Use(httpmw.RateLimit(1, time.Second))
 		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
 			rw.WriteHeader(http.StatusOK)
 		})
 
-		require.Eventually(t, func() bool {
+		for i := 0; i < 5; i++ {
 			req := httptest.NewRequest("GET", "/", nil)
 			rec := httptest.NewRecorder()
 			rtr.ServeHTTP(rec, req)
 			resp := rec.Result()
-			defer resp.Body.Close()
-			return resp.StatusCode == http.StatusTooManyRequests
-		}, testutil.WaitShort, testutil.IntervalFast)
+			_ = resp.Body.Close()
+			require.Equal(t, i != 0, resp.StatusCode == http.StatusTooManyRequests)
+		}
 	})
 
 	t.Run("RandomIPs", func(t *testing.T) {
 		t.Parallel()
 		rtr := chi.NewRouter()
-		rtr.Use(httpmw.RateLimit(5, time.Second))
+		// Because these are random IPs, the limit should never be hit!
+		rtr.Use(httpmw.RateLimit(1, time.Second))
 		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
 			rw.WriteHeader(http.StatusOK)
 		})
 
-		require.Never(t, func() bool {
+		for i := 0; i < 5; i++ {
 			req := httptest.NewRequest("GET", "/", nil)
 			rec := httptest.NewRecorder()
 			req.RemoteAddr = randRemoteAddr()
 			rtr.ServeHTTP(rec, req)
 			resp := rec.Result()
-			defer resp.Body.Close()
-			return resp.StatusCode == http.StatusTooManyRequests
-		}, testutil.WaitShort, testutil.IntervalFast)
+			_ = resp.Body.Close()
+			require.False(t, resp.StatusCode == http.StatusTooManyRequests)
+		}
 	})
 
 	t.Run("RegularUser", func(t *testing.T) {
@@ -81,7 +81,7 @@ func TestRateLimit(t *testing.T) {
 			Optional: false,
 		}))
 
-		rtr.Use(httpmw.RateLimit(5, time.Second))
+		rtr.Use(httpmw.RateLimit(1, time.Second))
 		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
 			rw.WriteHeader(http.StatusOK)
 		})
@@ -98,17 +98,17 @@ func TestRateLimit(t *testing.T) {
 		defer resp.Body.Close()
 		require.Equal(t, http.StatusPreconditionRequired, resp.StatusCode)
 
-		require.Eventually(t, func() bool {
+		for i := 0; i < 5; i++ {
 			req := httptest.NewRequest("GET", "/", nil)
 			req.Header.Set(codersdk.SessionTokenHeader, key)
 			rec := httptest.NewRecorder()
 			// Assert we're not using IP address.
 			req.RemoteAddr = randRemoteAddr()
 			rtr.ServeHTTP(rec, req)
 			resp := rec.Result()
-			defer resp.Body.Close()
-			return resp.StatusCode == http.StatusTooManyRequests
-		}, testutil.WaitShort, testutil.IntervalFast)
+			_ = resp.Body.Close()
+			require.Equal(t, i != 0, resp.StatusCode == http.StatusTooManyRequests)
+		}
 	})
 
 	t.Run("OwnerBypass", func(t *testing.T) {
@@ -127,12 +127,12 @@ func TestRateLimit(t *testing.T) {
 			Optional: false,
 		}))
 
-		rtr.Use(httpmw.RateLimit(5, time.Second))
+		rtr.Use(httpmw.RateLimit(1, time.Second))
 		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
 			rw.WriteHeader(http.StatusOK)
 		})
 
-		require.Never(t, func() bool {
+		for i := 0; i < 5; i++ {
 			req := httptest.NewRequest("GET", "/", nil)
 			req.Header.Set(codersdk.SessionTokenHeader, key)
 			req.Header.Set(codersdk.BypassRatelimitHeader, "true")
@@ -141,8 +141,8 @@ func TestRateLimit(t *testing.T) {
 			req.RemoteAddr = randRemoteAddr()
 			rtr.ServeHTTP(rec, req)
 			resp := rec.Result()
-			defer resp.Body.Close()
-			return resp.StatusCode == http.StatusTooManyRequests
-		}, testutil.WaitShort, testutil.IntervalFast)
+			_ = resp.Body.Close()
+			require.False(t, resp.StatusCode == http.StatusTooManyRequests)
+		}
 	})
 }
diff --git a/provisioner/terraform/install_test.go b/provisioner/terraform/install_test.go
@@ -18,6 +18,9 @@ import (
 
 func TestInstall(t *testing.T) {
 	t.Parallel()
+	if testing.Short() {
+		t.SkipNow()
+	}
 	ctx := context.Background()
 	dir := t.TempDir()
 	log := slogtest.Make(t, nil)
