TestRateLimitByFullPath

Zondax · Dec 24, 2023 · 1d06de7 · 1d06de7
1 parent e569109
commit 1d06de7
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/pkg/zrouter/zmiddlewares/rate_limit_test.go b/pkg/zrouter/zmiddlewares/rate_limit_test.go
@@ -28,3 +28,24 @@ func TestRateLimit(t *testing.T) {
 	r.ServeHTTP(rec, req)
 	assert.Equal(t, http.StatusTooManyRequests, rec.Code)
 }
+
+func TestRateLimitByFullPath(t *testing.T) {
+	r := chi.NewRouter()
+
+	r.Use(RateLimitByFullPath(1))
+
+	r.Get("/tests/{id}", func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("OK"))
+	})
+
+	req := httptest.NewRequest("GET", "/tests/1", nil)
+
+	rec := httptest.NewRecorder()
+	r.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "OK", rec.Body.String())
+
+	rec = httptest.NewRecorder()
+	r.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusTooManyRequests, rec.Code)
+}