ModelRouter/internal/gateway/key_quota_test.go

package gateway

import (
	"testing"
	"time"

	"llmsproxy/internal/config"
)

// ---- data layer: per-key isolation ----

func TestKeyWindowTokensIsolatesKeys(t *testing.T) {
	s := NewStats(100)
	now := time.Now().UnixMilli()
	s.Record(Req{Time: now, Key: "keyA", Model: "m1", Prompt: 500, Compl: 500, OK: true, Status: 200})
	s.Record(Req{Time: now, Key: "keyB", Model: "m1", Prompt: 500, Compl: 500, OK: true, Status: 200})

	if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 1000 {
		t.Errorf("keyA model tokens = %d, want 1000 (its own usage only)", got)
	}
	if got := s.KeyWindowModelTokens("keyB", "m1", "", 3600); got != 1000 {
		t.Errorf("keyB model tokens = %d, want 1000", got)
	}
	// the key-blind bucket stays global on purpose: it backs the AUTO slot
	// quota, which limits the whole gateway, not one key.
	if got := s.WindowTokens("m1", "", 3600); got != 2000 {
		t.Errorf("WindowTokens = %d, want 2000 (global per-model total must be unchanged)", got)
	}
}

// A source pin scopes the cap to that one upstream. The pinned bucket is
// maintained lazily (it is only read by quotas that actually pin a source), so
// usage recorded before anything queried the pin is not attributed to it.
func TestKeyWindowTokensSeparatesSourcePin(t *testing.T) {
	s := NewStats(100)
	now := time.Now().UnixMilli()
	s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200}) // 200 tok
	s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 200, Compl: 200, OK: true, Status: 200}) // 400 tok

	// The unpinned bucket always counts — it is what a scope entry without a
	// source pin reads. 200 + 400.
	if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 600 {
		t.Errorf("unpinned model tokens = %d, want 600 (both sources)", got)
	}

	// Reading a pin opts this (key, model) into pinned accounting.
	if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 0 {
		t.Errorf("srcX pin before opt-in = %d, want 0 (the lazy bucket has not accrued yet)", got)
	}

	// From now on both pins accrue.
	more := now + 1
	s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 10, Compl: 10, OK: true, Status: 200}) // 20 tok
	s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 30, Compl: 30, OK: true, Status: 200}) // 60 tok
	if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 20 {
		t.Errorf("srcX-pinned tokens after opt-in = %d, want 20", got)
	}
	if got := s.KeyWindowModelTokens("keyA", "m1", "srcY", 3600); got != 60 {
		t.Errorf("srcY-pinned tokens = %d, want 60 (a pin read for one source must not blind the other)", got)
	}
}

// This is the shape of every real chat record: Source is always populated.
// The unpinned bucket must still see the tokens, or a per-model quota without
// a source pin reads an empty bucket and never trips.
func TestKeyWindowModelTokensUnpinnedSeesSourcedTraffic(t *testing.T) {
	s := NewStats(100)
	s.Record(Req{Time: time.Now().UnixMilli(), Key: "keyA", Model: "m1", Source: "deepseek",
		Prompt: 1000, Compl: 1000, OK: true, Status: 200})

	if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 2000 {
		t.Errorf("unpinned tokens = %d, want 2000 — an unpinned per-model quota would never trip otherwise", got)
	}
}

func TestKeyWindowRespectsResetWindow(t *testing.T) {
	s := NewStats(100)
	old := time.Now().Add(-30 * 24 * time.Hour).UnixMilli()
	s.Record(Req{Time: old, Key: "keyA", Model: "m1", Prompt: 100, Compl: 100, OK: true, Status: 200})
	s.Record(Req{Time: time.Now().UnixMilli(), Key: "keyA", Model: "m1", Prompt: 5, Compl: 5, OK: true, Status: 200})

	if got := s.KeyWindowTokens("keyA", 3600); got != 10 {
		t.Errorf("1h window = %d, want 10 (30-day-old usage must not count)", got)
	}
	if got := s.KeyWindowTokens("keyA", 0); got != 210 {
		t.Errorf("all-time total = %d, want 210", got)
	}
}

func TestKeyWindowReqsCountsEveryRequest(t *testing.T) {
	s := NewStats(100)
	now := time.Now().UnixMilli()
	// a failed request and an image request both count: a client that loops on
	// failures must still burn its request quota
	s.Record(Req{Time: now, Key: "keyA", Model: "m1", Type: "chat", OK: true, Status: 200})
	s.Record(Req{Time: now, Key: "keyA", Model: "m1", Type: "chat", OK: false, Status: 500})
	s.Record(Req{Time: now, Key: "keyA", Model: "img", Type: "image", OK: true, Status: 200})
	s.Record(Req{Time: now, Key: "keyB", Model: "m1", Type: "chat", OK: true, Status: 200})

	if got := s.KeyWindowReqs("keyA", 3600); got != 3 {
		t.Errorf("keyA reqs = %d, want 3 (chat ok + chat fail + image)", got)
	}
	if got := s.KeyWindowReqs("keyB", 3600); got != 1 {
		t.Errorf("keyB reqs = %d, want 1", got)
	}
}

func TestKeyWindowTokensZeroTokensNotCounted(t *testing.T) {
	s := NewStats(100)
	now := time.Now().UnixMilli()
	// a request that reported no usage must not create a bucket entry
	s.Record(Req{Time: now, Key: "keyA", Model: "", Type: "chat", OK: true, Status: 200})
	if got := s.KeyWindowTokens("keyA", 3600); got != 0 {
		t.Errorf("tokens = %d, want 0", got)
	}
	if got := s.KeyWindowReqs("keyA", 3600); got != 1 {
		t.Errorf("reqs = %d, want 1 (request still happened)", got)
	}
}

// ---- reset window arithmetic ----

func TestAutoSecondsToReset(t *testing.T) {
	if got := AutoSecondsToReset("", 0); got != 0 {
		t.Errorf("no period = %d, want 0 (never resets -> no retry hint)", got)
	}
	now := time.Now().Unix()
	for _, tc := range []struct {
		period string
		hours  int64
		want   int64
	}{
		{"hour", 0, 3600},
		{"week", 0, 7 * 24 * 3600},
		{"month", 0, 30 * 24 * 3600},
		{"nhour", 6, 6 * 3600},
	} {
		got := AutoSecondsToReset(tc.period, tc.hours)
		if got <= 0 || got > tc.want {
			t.Errorf("AutoSecondsToReset(%q,%d) = %d, want in (0,%d]", tc.period, tc.hours, got, tc.want)
		}
		// must never exceed the window itself
		if got < tc.want-now%3600 {
			t.Logf("note: %q hint %ds < remaining-window %ds (rounds to hour boundary)", tc.period, got, tc.want-now%3600)
		}
		if got > tc.want {
			t.Errorf("hint %d exceeds window %d", got, tc.want)
		}
	}
	_ = now
}

func TestAutoSecondsToResetNeverExceedsWindow(t *testing.T) {
	// "nhour" with a tiny window must not hand out a longer wait than the
	// window itself (which would stall a client past its own reset)
	for hours := int64(1); hours <= 48; hours++ {
		got := AutoSecondsToReset("nhour", hours)
		if got <= 0 || got > hours*3600 {
			t.Errorf("nhour/%d = %d, want in (0,%d]", hours, got, hours*3600)
		}
	}
}

// ---- config validation ----

func TestKeyQuotaValidate(t *testing.T) {
	cases := []struct {
		name    string
		q       config.KeyQuota
		wantErr bool
	}{
		{"unlimited", config.KeyQuota{}, false},
		{"tokens hourly", config.KeyQuota{TokenQuota: 1000, Period: "hour"}, false},
		{"tokens n-hour", config.KeyQuota{TokenQuota: 1000, Period: "nhour", Hours: 6}, false},
		{"reqs weekly", config.KeyQuota{ReqQuota: 100, Period: "week"}, false},
		{"no period never resets", config.KeyQuota{TokenQuota: 1000}, false},
		{"negative tokens", config.KeyQuota{TokenQuota: -1}, true},
		{"negative reqs", config.KeyQuota{ReqQuota: -1}, true},
		{"negative hours", config.KeyQuota{TokenQuota: 5, Hours: -1}, true},
		{"typo period becomes all-time if accepted", config.KeyQuota{TokenQuota: 5, Period: "houre"}, true},
		{"nhour without hours", config.KeyQuota{TokenQuota: 5, Period: "nhour"}, true},
		{"nhour with 0 hours", config.KeyQuota{TokenQuota: 5, Period: "nhour", Hours: 0}, true},
		{"junk period with only req quota", config.KeyQuota{ReqQuota: 5, Period: "daily"}, true},
		{"junk period with no caps is irrelevant", config.KeyQuota{Period: "daily"}, false},
	}
	for _, tc := range cases {
		err := tc.q.Validate()
		if tc.wantErr && err == nil {
			t.Errorf("%s: expected error, got nil", tc.name)
		}
		if !tc.wantErr && err != nil {
			t.Errorf("%s: unexpected error: %v", tc.name, err)
		}
	}
}

func TestNormalizeRole(t *testing.T) {
	if got := config.NormalizeRole(""); got != "user" {
		t.Errorf("empty role = %q, want user", got)
	}
	if got := config.NormalizeRole("admin"); got != "admin" {
		t.Errorf("admin role = %q, want admin", got)
	}
	if got := config.NormalizeRole("root"); got != "user" {
		t.Errorf("unknown role = %q, want user (never escalate)", got)
	}
}

// Retention must hold for the lazy pinned buckets too, or a long-lived key
// would grow without bound.
func TestPinnedBucketRespectsRetention(t *testing.T) {
	s := NewStats(100)
	_ = s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600) // opt in
	nowH := time.Now().Unix() / 3600
	for h := int64(0); h < quotaRetentionHours+50; h++ {
		s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyA", Model: "m1", Source: "srcX",
			Prompt: 10, Compl: 10, OK: true, Status: 200})
	}
	s.mu.Lock()
	n := len(s.keyModelHour["keyA"]["srcX::m1"])
	s.mu.Unlock()
	if n > quotaRetentionHours {
		t.Errorf("pinned bucket holds %d hours, want <= %d", n, quotaRetentionHours)
	}
}