starter code

2026-09-06 15:15:21 +02:00
parent 445acb18da
commit 1aa524fb22
10 changed files with 1094 additions and 0 deletions
@@ -0,0 +1,208 @@
+package ratelimit
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestMemoryWindow(t *testing.T) {
+	now := time.Date(2026, 9, 4, 12, 0, 0, 0, time.UTC)
+	m := NewMemory()
+	m.Now = func() time.Time { return now }
+
+	ctx := context.Background()
+
+	for i := 1; i <= 3; i++ {
+		res, err := m.Allow(ctx, "k", 3, time.Minute)
+		if err != nil || !res.Allowed {
+			t.Fatalf("request %d: allowed=%v err=%v", i, res.Allowed, err)
+		}
+		if res.Remaining != 3-i {
+			t.Errorf("request %d: remaining %d, want %d", i, res.Remaining, 3-i)
+		}
+	}
+
+	res, _ := m.Allow(ctx, "k", 3, time.Minute)
+	if res.Allowed {
+		t.Error("4th request allowed past a limit of 3")
+	}
+	if res.RetryAfter <= 0 {
+		t.Error("RetryAfter not set on a rejection")
+	}
+
+	// Window rolls over.
+	now = now.Add(time.Minute + time.Second)
+	if res, _ := m.Allow(ctx, "k", 3, time.Minute); !res.Allowed {
+		t.Error("request rejected after the window expired")
+	}
+}
+
+func TestMemoryKeysAreIndependent(t *testing.T) {
+	m := NewMemory()
+	ctx := context.Background()
+
+	for i := 0; i < 5; i++ {
+		m.Allow(ctx, "a", 5, time.Minute)
+	}
+
+	if res, _ := m.Allow(ctx, "a", 5, time.Minute); res.Allowed {
+		t.Error("key a not limited")
+	}
+	if res, _ := m.Allow(ctx, "b", 5, time.Minute); !res.Allowed {
+		t.Error("key b affected by key a")
+	}
+}
+
+func TestMemoryZeroLimitIsUnlimited(t *testing.T) {
+	m := NewMemory()
+	for i := 0; i < 1000; i++ {
+		if res, _ := m.Allow(context.Background(), "k", 0, time.Minute); !res.Allowed {
+			t.Fatalf("zero limit rejected at %d", i)
+		}
+	}
+}
+
+func TestMemorySweep(t *testing.T) {
+	now := time.Now()
+	m := NewMemory()
+	m.Now = func() time.Time { return now }
+
+	m.Allow(context.Background(), "a", 5, time.Minute)
+	m.Allow(context.Background(), "b", 5, time.Hour)
+
+	if n := m.Sweep(); n != 0 {
+		t.Errorf("swept %d live windows", n)
+	}
+
+	now = now.Add(2 * time.Minute)
+
+	if n := m.Sweep(); n != 1 {
+		t.Errorf("swept %d, want 1 (a expired, b did not)", n)
+	}
+}
+
+func TestMemoryIsConcurrencySafe(t *testing.T) {
+	m := NewMemory()
+	ctx := context.Background()
+
+	const goroutines, each = 20, 50
+
+	done := make(chan int, goroutines)
+
+	for g := 0; g < goroutines; g++ {
+		go func() {
+			allowed := 0
+			for i := 0; i < each; i++ {
+				if res, _ := m.Allow(ctx, "shared", 100, time.Minute); res.Allowed {
+					allowed++
+				}
+			}
+			done <- allowed
+		}()
+	}
+
+	total := 0
+	for g := 0; g < goroutines; g++ {
+		total += <-done
+	}
+
+	// Exactly 100 of the 1000 attempts may pass. A check-then-increment race
+	// would let more through.
+	if total != 100 {
+		t.Errorf("%d requests allowed under concurrency, want exactly 100", total)
+	}
+}
+
+// fakeRunner stands in for Redis so the script result handling can be tested.
+type fakeRunner struct {
+	ret  any
+	err  error
+	keys []string
+	args []any
+}
+
+func (f *fakeRunner) Eval(_ context.Context, _ string, keys []string, args ...any) (any, error) {
+	f.keys, f.args = keys, args
+	return f.ret, f.err
+}
+
+func TestRedisResultParsing(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("under limit", func(t *testing.T) {
+		f := &fakeRunner{ret: []any{int64(3), int64(45000)}}
+		r := NewRedis(f)
+
+		res, err := r.Allow(ctx, "key:1", 10, time.Minute)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !res.Allowed || res.Remaining != 7 {
+			t.Errorf("allowed=%v remaining=%d, want true/7", res.Allowed, res.Remaining)
+		}
+		if res.RetryAfter != 45*time.Second {
+			t.Errorf("RetryAfter %v, want 45s", res.RetryAfter)
+		}
+		if len(f.keys) != 1 || f.keys[0] != "rl:key:1" {
+			t.Errorf("keys %v, want [rl:key:1]", f.keys)
+		}
+	})
+
+	t.Run("over limit", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{int64(11), int64(1000)}})
+		res, _ := r.Allow(ctx, "k", 10, time.Minute)
+		if res.Allowed || res.Remaining != 0 {
+			t.Errorf("allowed=%v remaining=%d, want false/0", res.Allowed, res.Remaining)
+		}
+	})
+
+	t.Run("exactly at limit is allowed", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{int64(10), int64(1000)}})
+		if res, _ := r.Allow(ctx, "k", 10, time.Minute); !res.Allowed {
+			t.Error("request at exactly the limit was rejected")
+		}
+	})
+
+	t.Run("plain int accepted", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{1, 1000}})
+		if _, err := r.Allow(ctx, "k", 10, time.Minute); err != nil {
+			t.Errorf("int result rejected: %v", err)
+		}
+	})
+
+	t.Run("missing ttl falls back to the window", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: []any{int64(11), int64(-1)}})
+		res, _ := r.Allow(ctx, "k", 10, time.Minute)
+		if res.RetryAfter != time.Minute {
+			t.Errorf("RetryAfter %v, want 1m", res.RetryAfter)
+		}
+	})
+
+	t.Run("eval error propagates", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{err: errors.New("nope")})
+		if _, err := r.Allow(ctx, "k", 10, time.Minute); err == nil {
+			t.Error("expected an error so the caller can fail open")
+		}
+	})
+
+	t.Run("garbage result is an error not a panic", func(t *testing.T) {
+		r := NewRedis(&fakeRunner{ret: "not a list"})
+		if _, err := r.Allow(ctx, "k", 10, time.Minute); err == nil {
+			t.Error("expected an error")
+		}
+	})
+
+	t.Run("zero limit skips redis entirely", func(t *testing.T) {
+		f := &fakeRunner{ret: []any{int64(1), int64(1)}}
+		r := NewRedis(f)
+		res, _ := r.Allow(ctx, "k", 0, time.Minute)
+		if !res.Allowed {
+			t.Error("zero limit rejected")
+		}
+		if f.keys != nil {
+			t.Error("redis was called for an unlimited tier")
+		}
+	})
+}