Short and single-letter names in application code become full words (tests keep theirs). Wrapped calls, signatures and composite literals put one item per line, and one-line function bodies are expanded. Exported names and groups of constants get doc comments, and comments refer to the renamed identifiers.
141 lines
3.8 KiB
Go
141 lines
3.8 KiB
Go
package mock
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRateLimitWindows(t *testing.T) {
|
|
now := time.Date(2026, 10, 6, 10, 3, 20, 0, time.UTC) // first request of the hour
|
|
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")
|
|
}
|
|
}
|