package test import ( "RenewCMS/api/pkg/apperr" "net/http" "net/http/httptest" "testing" ) const publicReadBudget = 120 func (e *locPublicEnv) getFrom(remoteAddr string, path string) *httptest.ResponseRecorder { req := httptest.NewRequest(http.MethodGet, path, nil) req.RemoteAddr = remoteAddr rec := httptest.NewRecorder() e.router.ServeHTTP(rec, req) return rec } func TestPublicReadRateLimit(t *testing.T) { t.Run("The public article routes and the sitemap share one budget per client", func(t *testing.T) { e := newLocPublicEnv(t) client := "198.51.100.7:4000" for index := 0; index < publicReadBudget; index++ { path := locFrBase if index%2 == 0 { path = "/v1/site/journal/sitemap" } if rec := e.getFrom(client, path); rec.Code != http.StatusOK { t.Fatalf("request %d: %d %s", index, rec.Code, rec.Body.String()) } } for _, path := range []string{locFrBase, "/v1/site/journal/sitemap", "/v1/public/site/journal/story/5"} { rec := e.getFrom(client, path) if rec.Code != http.StatusTooManyRequests || rec.Body.String() != apperr.TooManyRequestsErrorMessage+"\n" || rec.Header().Get("Retry-After") == "" { t.Errorf("%s: %d %q", path, rec.Code, rec.Body.String()) } } if rec := e.getFrom("198.51.100.8:4000", locFrBase); rec.Code != http.StatusOK { t.Errorf("another client is limited: %d", rec.Code) } }) }