package mock import ( "net/http" "net/http/httptest" "strconv" "strings" "testing" "time" ) func TestRateLimitWindows(t *testing.T) { // The first request of the hour. now := time.Date(2026, 10, 6, 10, 3, 20, 0, time.UTC) rl := NewRateLimiter(12, 5*time.Minute) rl.now = func() time.Time { return now } for i := 0; i < 12; i++ { if ok, _ := rl.allow("a"); !ok { t.Fatalf("request %d rejected", i+1) } } now = now.Add(100 * time.Second) // 10:05:00, window ends 10:08:20 if ok, secs := rl.allow("a"); ok || secs != 200 { t.Fatalf("want reject with 200s, got ok=%v secs=%d", ok, secs) } if ok, _ := rl.allow("b"); !ok { t.Fatal("other client should be unaffected") } now = time.Date(2026, 10, 6, 10, 8, 20, 0, time.UTC) // next window opens if ok, _ := rl.allow("a"); !ok { t.Fatal("new window should allow") } // Window 10:58:20-11:03:20 is cut short by the new hour at 11:00:00. now = time.Date(2026, 10, 6, 10, 58, 30, 0, time.UTC) for i := 0; i < 12; i++ { rl.allow("a") } if ok, secs := rl.allow("a"); ok || secs != 90 { t.Fatalf("want reject with 90s, got ok=%v secs=%d", ok, secs) } now = time.Date(2026, 10, 6, 11, 0, 5, 0, time.UTC) // re-anchors here if ok, _ := rl.allow("a"); !ok { t.Fatal("new hour should allow") } } func TestClientKey(t *testing.T) { cases := map[string]string{ "192.0.2.1:1234": "192.0.2.1", "[::1]:8080": "::1", "no-port": "no-port", // not host:port, used as-is } for addr, want := range cases { req := httptest.NewRequest("GET", "/", nil) req.RemoteAddr = addr if got := clientKey(req); got != want { t.Errorf("clientKey(%q) = %q, want %q", addr, got, want) } } } func TestRateLimitMiddleware(t *testing.T) { now := time.Date(2026, 10, 6, 10, 0, 0, 0, time.UTC) rl := NewRateLimiter(2, 5*time.Minute) rl.now = func() time.Time { return now } calls := 0 h := rl.Middleware(func(w http.ResponseWriter, req *http.Request) { calls++ w.WriteHeader(http.StatusOK) }) do := func(addr string) *httptest.ResponseRecorder { req := httptest.NewRequest("GET", "/", nil) req.RemoteAddr = addr rec := httptest.NewRecorder() h(rec, req) return rec } for i := 0; i < 2; i++ { if rec := do("192.0.2.1:1000"); rec.Code != http.StatusOK { t.Fatalf("request %d: status %d", i+1, rec.Code) } } // Same host, different port: still the same client. rec := do("192.0.2.1:2000") if rec.Code != http.StatusTooManyRequests { t.Fatalf("status %d, want 429", rec.Code) } if got := rec.Header().Get("Retry-After"); got != "300" { t.Errorf("Retry-After = %q, want 300", got) } if !strings.Contains(rec.Body.String(), "Rate limit of 2 exceeded") || !strings.Contains(rec.Body.String(), "300 seconds") { t.Errorf("body = %q", rec.Body.String()) } if calls != 2 { t.Errorf("next called %d times, want 2", calls) } if rec := do("192.0.2.2:1000"); rec.Code != http.StatusOK { t.Errorf("other client: status %d", rec.Code) } } // Once more than 10000 clients are tracked, a new client in a new hour drops // every client left over from earlier hours. func TestRateLimitPrunesStaleClients(t *testing.T) { now := time.Date(2026, 10, 6, 10, 30, 0, 0, time.UTC) rl := NewRateLimiter(1, time.Minute) rl.now = func() time.Time { return now } for i := 0; i <= 10000; i++ { rl.allow(strconv.Itoa(i)) } if len(rl.clients) != 10001 { t.Fatalf("tracking %d clients", len(rl.clients)) } // Same hour: nothing is stale, so nothing is dropped. rl.allow("same-hour") if len(rl.clients) != 10002 { t.Fatalf("same hour: tracking %d clients, want 10002", len(rl.clients)) } now = now.Add(time.Hour) rl.allow("new") if len(rl.clients) != 1 { t.Fatalf("after prune: tracking %d clients, want 1", len(rl.clients)) } if _, ok := rl.clients["new"]; !ok { t.Fatal("new client was pruned") } }