NextMyDay/internal/core/middleware/middleware_test.go

package middleware_test

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"git.lohmar.co.uk/cclohmar/NextMyDay/internal/core/middleware"
)

func TestNewCorrelationID(t *testing.T) {
	a := middleware.NewCorrelationID()
	b := middleware.NewCorrelationID()
	if a == "" || a == b {
		t.Fatalf("correlation ids are empty or identical: %q, %q", a, b)
	}
	if len(a) != 36 || a[8] != '-' || a[13] != '-' || a[18] != '-' || a[23] != '-' {
		t.Errorf("correlation id %q is not UUID-shaped", a)
	}
}

func TestCorrelationIDMiddleware(t *testing.T) {
	var seen string
	h := middleware.CorrelationID(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seen = middleware.CorrelationFrom(r.Context())
		w.WriteHeader(http.StatusOK)
	}))

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest("GET", "/home", nil))

	header := rec.Header().Get(middleware.CorrelationHeader)
	if header == "" {
		t.Fatal("response is missing the correlation header")
	}
	if seen != header {
		t.Errorf("context correlation = %q, header = %q; want equal", seen, header)
	}
}

func TestCorrelationRoundTrip(t *testing.T) {
	if got := middleware.CorrelationFrom(context.Background()); got != "" {
		t.Errorf("CorrelationFrom(background) = %q, want empty", got)
	}
	ctx := middleware.WithCorrelation(context.Background(), "corr-1")
	if got := middleware.CorrelationFrom(ctx); got != "corr-1" {
		t.Errorf("CorrelationFrom = %q, want corr-1", got)
	}
}

func TestRateLimiterAllow(t *testing.T) {
	l := middleware.NewRateLimiter(2, time.Minute)
	if !l.Allow("k") || !l.Allow("k") {
		t.Fatal("the first two requests should be allowed")
	}
	if l.Allow("k") {
		t.Error("the third request should be denied")
	}
	if !l.Allow("other") {
		t.Error("a different key must have its own bucket")
	}
}

func TestRateLimiterWindowResets(t *testing.T) {
	l := middleware.NewRateLimiter(1, time.Minute)
	now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
	l.SetClock(func() time.Time { return now })

	if !l.Allow("k") || l.Allow("k") {
		t.Fatal("limit of 1 not enforced")
	}
	now = now.Add(2 * time.Minute)
	if !l.Allow("k") {
		t.Error("the window did not reset")
	}
}

// TestRateLimiterPrunesElapsedWindows covers M3s: the state map must not grow
// without bound. Once enough distinct keys have been seen, an elapsed window
// makes the next insert sweep them away.
func TestRateLimiterPrunesElapsedWindows(t *testing.T) {
	l := middleware.NewRateLimiter(1, time.Minute)
	now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
	l.SetClock(func() time.Time { return now })

	// Seed enough distinct keys to make a sweep worthwhile.
	const seeded = 5000
	for i := 0; i < seeded; i++ {
		l.Allow("k" + strconv.Itoa(i))
	}
	if got := l.Len(); got != seeded {
		t.Fatalf("Len = %d, want %d live windows", got, seeded)
	}

	// Two windows later every seeded window has elapsed; the next new key
	// triggers the sweep.
	now = now.Add(2 * time.Minute)
	l.Allow("fresh")

	if got := l.Len(); got > 10 {
		t.Errorf("Len after the sweeps = %d, want the %d elapsed windows pruned (M3s)", got, seeded)
	}
}

// TestRateLimiterMapIsCapped covers M3s's hard cap: a flood of distinct keys
// must not grow the map without bound.
func TestRateLimiterMapIsCapped(t *testing.T) {
	l := middleware.NewRateLimiter(1, time.Minute)
	// The clock is frozen, so the time-gated sweep cannot fire — only the hard
	// cap bounds the map.
	now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
	l.SetClock(func() time.Time { return now })

	const inserted = 70_000
	for i := 0; i < inserted; i++ {
		l.Allow("k" + strconv.Itoa(i))
	}
	if got := l.Len(); got >= inserted {
		t.Errorf("the limiter held %d keys after %d inserts — the map is unbounded (M3s)", got, inserted)
	}
}

func TestRateLimiterMiddleware429(t *testing.T) {
	l := middleware.NewRateLimiter(1, time.Minute)
	h := l.Middleware(func(*http.Request) string { return "k" },
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }))

	rec1 := httptest.NewRecorder()
	h.ServeHTTP(rec1, httptest.NewRequest("GET", "/access", nil))
	if rec1.Code != http.StatusOK {
		t.Fatalf("first request = %d, want 200", rec1.Code)
	}

	rec2 := httptest.NewRecorder()
	h.ServeHTTP(rec2, httptest.NewRequest("GET", "/access", nil))
	if rec2.Code != http.StatusTooManyRequests {
		t.Fatalf("second request = %d, want 429", rec2.Code)
	}
	if rec2.Header().Get("Retry-After") == "" {
		t.Error("429 is missing a Retry-After header")
	}
}

func TestClassAndAuthSensitive(t *testing.T) {
	cases := []struct {
		path string
		want string
	}{
		{"/", ""},
		{"/home", ""},
		{"/api/v1", ""},
		{"/static/x.css", ""},
		{"/access", "access"},
		{"/setup", "setup"},
		{"/AbCdEf0123456789", "guid"},
		{"/AbCdEf0123456789/home", "guid"},
	}
	for _, tc := range cases {
		if got := middleware.Class(tc.path); got != tc.want {
			t.Errorf("Class(%q) = %q, want %q", tc.path, got, tc.want)
		}
		if got := middleware.AuthSensitive(tc.path); got != (tc.want != "") {
			t.Errorf("AuthSensitive(%q) = %v, want %v", tc.path, got, tc.want != "")
		}
	}
}