package quota import ( "context" "fmt" "net/http" "net/http/httptest" "strings" "testing" "time" "ukrrs.com/mopac/harness/internal/config" ) func TestParseUsage(t *testing.T) { ok := `{"usage":{"five_hour":{"used_credits":15400,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"}, "weekly":{"used_credits":130200,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}` alias := `{"usage":{"five_hour":{"credits_used":100,"limit_credits":200,"reset_time":"2026-08-29T22:00:00Z"}, "weekly":{"credits_used":"150","limit_credits":"200","reset_time":"2026-09-02T00:00:00Z"}}}` bad := []struct { name, body string }{ {"html error page", `502 Bad Gateway`}, {"missing buckets", `{"usage":{}}`}, {"missing used", `{"usage":{"five_hour":{"total_credits":28000},"weekly":{"used_credits":1,"total_credits":2}}}`}, {"bad reset", `{"usage":{"five_hour":{"used_credits":1,"total_credits":2,"reset_at":"tomorrow"},"weekly":{"used_credits":1,"total_credits":2}}}`}, } cases := []struct { name string body string want5h float64 want5p float64 wantWkd float64 wantErr bool }{ {"canonical shape", ok, 15400, 55.0, 93.0, false}, {"field aliases + string numbers", alias, 100, 50, 75, false}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { snap, err := ParseUsage("zai-1", time.Now(), []byte(c.body)) if err != nil { t.Fatalf("ParseUsage: %v", err) } b5, ok := snap.Bucket(Bucket5h) if !ok || b5.Used != c.want5h || !nearly(b5.UsedPct(), c.want5p) { t.Errorf("5h bucket = %+v, want used=%v pct=%v", b5, c.want5h, c.want5p) } bw, _ := snap.Bucket(BucketWeekly) if !nearly(bw.UsedPct(), c.wantWkd) { t.Errorf("weekly pct = %v, want %v", bw.UsedPct(), c.wantWkd) } if got := snap.MaxUsedPct(); !nearly(got, c.wantWkd) { t.Errorf("MaxUsedPct = %v, want %v", got, c.wantWkd) } }) } for _, c := range bad { t.Run(c.name, func(t *testing.T) { if _, err := ParseUsage("zai-1", time.Now(), []byte(c.body)); err == nil { t.Fatal("expected error, got nil (a garbage body must never read as 'quota fine')") } }) } } func nearly(a, b float64) bool { d := a - b return d < 1e-9 && d > -1e-9 } func TestParseUsageZeroLimitNeverBlocks(t *testing.T) { snap, err := ParseUsage("a", time.Now(), []byte(`{"usage":{"five_hour":{"used_credits":999,"total_credits":0,"reset_at":""},"weekly":{"used_credits":1,"total_credits":2}}}`)) if err != nil { t.Fatal(err) } if b, _ := snap.Bucket(Bucket5h); b.UsedPct() != 0 { t.Errorf("zero-limit bucket must read 0%%, got %v", b.UsedPct()) } } func TestGateUsesBearerAndNeverLogsKey(t *testing.T) { var gotAuth string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuth = r.Header.Get("Authorization") fmt.Fprintf(w, `{"usage":{"five_hour":{"used_credits":1,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"},"weekly":{"used_credits":1,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}`) })) defer srv.Close() cfg := gateTestCfg() cfg.UsageURL = srv.URL cfg.KeyRef = "literal:seekrit-zai-key" gate := newGate(t, cfg) snap := gate.Snapshot(context.Background()) if snap.Source != "usage_url" { t.Fatalf("source = %s, want usage_url", snap.Source) } if gotAuth != "Bearer seekrit-zai-key" { t.Errorf("auth header = %q", gotAuth) } // Poll failures (5xx) must fall back without panicking and never leak // the key into the returned error strings. srv.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "boom", http.StatusInternalServerError) }) gate2 := newGate(t, cfg) snap2 := gate2.Snapshot(context.Background()) if snap2 == nil { t.Fatal("nil snapshot on poll failure; gate must degrade to estimates") } if strings.Contains(snap2.Source, "seekrit") { t.Errorf("key leaked into snapshot source %q", snap2.Source) } } func TestGateSharedSnapshotRoundTrip(t *testing.T) { fake := newFakeRedis(t) st, err := NewSharedState(fake.url(t)) if err != nil { t.Fatal(err) } defer st.Close() snap := &QuotaSnapshot{Account: "zai-1", Source: "usage_url", FetchedAt: time.Now(), Buckets: []Bucket{{ID: Bucket5h, Used: 1, Limit: 28000}}} if err := st.PublishSnapshot(snap, time.Minute); err != nil { t.Fatal(err) } got, err := st.LoadSnapshot("zai-1") if err != nil || got == nil || got.Source != "usage_url" { t.Fatalf("round trip: %v %+v", err, got) } if b, _ := got.Bucket(Bucket5h); b.Used != 1 { t.Errorf("bucket lost: %+v", got.Buckets) } } func TestEstimateCredits(t *testing.T) { cases := []struct { name string model string prompt, cached, comp int peak bool want float64 }{ // flagship: (10000*6.9 + 2000*1.7 + 10000*24)/10000 = 31.24 {"flagship peak", "glm-5.3", 10000, 2000, 10000, true, 31.24}, // off-peak halves: 15.62 {"flagship off-peak", "glm-5.3", 10000, 2000, 10000, false, 15.62}, // flash: (10000*2.3 + 2000*0.56 + 10000*8)/10000 = 10.412 -> /2 {"flash off-peak", "glm-4.7-flash", 10000, 2000, 10000, false, 5.206}, {"flash peak", "GLM-5.3-Flash", 10000, 0, 10000, true, 10.3}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { got := EstimateCredits(c.model, c.prompt, c.cached, c.comp, c.peak) if diff := got - c.want; diff > 1e-9 || diff < -1e-9 { t.Errorf("credits = %v, want %v", got, c.want) } }) } if MultipliersFor("glm-9.9-omega") != FlagshipMultipliers() { t.Error("unknown model must read as flagship (conservative)") } } func TestScheduleEdges(t *testing.T) { cst, err := time.LoadLocation("America/Chicago") if err != nil { t.Fatal(err) } at := func(day string, hm string) time.Time { tm, err := time.ParseInLocation("2006-01-02 15:04", day+" "+hm, cst) if err != nil { t.Fatal(err) } return tm } // Default window 01:00-05:00 CST, weekdays only. sched, err := NewSchedule("01:00", "05:00", "America/Chicago", true) if err != nil { t.Fatal(err) } cases := []struct { name string t time.Time want bool }{ {"just before start (Fri)", at("2026-08-28", "00:59"), false}, {"at start (Fri)", at("2026-08-28", "01:00"), true}, {"mid window (Fri)", at("2026-08-28", "03:00"), true}, {"last minute (Fri)", at("2026-08-28", "04:59"), true}, {"at end (Fri)", at("2026-08-28", "05:00"), false}, {"evening off-peak (Fri)", at("2026-08-28", "19:00"), false}, {"weekend inside window (Sat)", at("2026-08-29", "02:00"), false}, {"weekend inside window (Sun)", at("2026-08-30", "02:00"), false}, {"Monday early window", at("2026-08-31", "02:00"), true}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { if got := sched.InPeak(c.t); got != c.want { t.Errorf("InPeak(%s) = %v, want %v", c.t, got, c.want) } }) } // Wrapped window 22:00-02:00, weekdays only: evening part carries the // start day's weekday; morning part checks the PREVIOUS day. wrap, err := NewSchedule("22:00", "02:00", "America/Chicago", true) if err != nil { t.Fatal(err) } wrapCases := []struct { name string t time.Time want bool }{ {"Fri 23:00 (started Fri)", at("2026-08-28", "23:00"), true}, {"Sat 01:00 (window started Fri)", at("2026-08-29", "01:00"), true}, {"Sat 23:00 (started Sat)", at("2026-08-29", "23:00"), false}, {"Sun 01:00 (window started Sat)", at("2026-08-30", "01:00"), false}, } for _, c := range wrapCases { t.Run("wrap: "+c.name, func(t *testing.T) { if got := wrap.InPeak(c.t); got != c.want { t.Errorf("InPeak(%s) = %v, want %v", c.t, got, c.want) } }) } // Every-day window ignores weekday flags. daily, _ := NewSchedule("01:00", "05:00", "America/Chicago", false) if !daily.InPeak(at("2026-08-29", "02:00")) { // Saturday t.Error("weekdays_only=false must peak on Saturday too") } // UTC window evaluated on the INSTANT: 21:00 UTC is outside, but the // same wall-clock digits in CST (21:00 CDT == 02:00 UTC) are inside. utc, _ := NewSchedule("01:00", "05:00", "UTC", false) utcT := time.Date(2026, 8, 28, 21, 0, 0, 0, time.UTC) if utc.InPeak(utcT) { t.Error("21:00 UTC must be outside the UTC window") } cstWall := time.Date(2026, 8, 28, 21, 0, 0, 0, cst) if !utc.InPeak(cstWall) { t.Error("21:00 CST (== 02:00 UTC) must be inside the UTC window") } if _, err := NewSchedule("25:00", "05:00", "UTC", false); err == nil { t.Error("25:00 must fail to parse") } if _, err := NewSchedule("01:00", "05:00", "Mars/Olympus", false); err == nil { t.Error("unknown timezone must fail") } } func TestRecordTurnFeedsEstimateSnapshot(t *testing.T) { fake := newFakeRedis(t) cfg := gateTestCfg() cfg.RedisURL = fake.url(t) gate := newGate(t, cfg) // Friday 03:00 CST: inside the default peak window. fri := time.Date(2026, 8, 28, 3, 0, 0, 0, mustLoc(t, "America/Chicago")) gate.SetClock(func() time.Time { return fri }) credits := gate.RecordTurn(EstimateTurnInput{Model: "glm-5.3", PromptTokens: 10000, CachedTokens: 2000, CompletionTokens: 10000, Peak: true}) if credits != 31.24 { t.Fatalf("credits = %v, want 31.24", credits) } snap := gate.estimateSnapshot(fri) b5, _ := snap.Bucket(Bucket5h) bw, _ := snap.Bucket(BucketWeekly) if b5.Used < 31.23 || b5.Used > 31.25 || bw.Used < 31.23 || bw.Used > 31.25 { t.Errorf("estimate buckets after one turn: 5h=%v weekly=%v", b5.Used, bw.Used) } if b5.Limit != cfg.Plan5hCredits || bw.Limit != cfg.PlanWeeklyCredits { t.Errorf("limits = %v/%v, want plan limits", b5.Limit, bw.Limit) } // Local-only mode (no redis) keeps its own accounting. local := newGate(t, gateTestCfg()) local.SetClock(func() time.Time { return fri }) local.RecordTurn(EstimateTurnInput{Model: "glm-4.7-flash", PromptTokens: 10000, CompletionTokens: 10000, Peak: false}) ls := local.estimateSnapshot(fri) lb, _ := ls.Bucket(Bucket5h) if lb.Used < 5.14 || lb.Used > 5.16 { t.Errorf("local estimate = %v, want ~5.15", lb.Used) } } func mustLoc(t *testing.T, name string) *time.Location { t.Helper() l, err := time.LoadLocation(name) if err != nil { t.Fatal(err) } return l } // TestDecideNineteenWall replays the 2026-08-28 19:00 quota wall: weekly // bucket exhausted at 19:00 CST — the gate must DEFER every class with a // surfaced reason instead of letting the turn die at the provider. func TestDecideNineteenWall(t *testing.T) { gate := newGate(t, gateTestCfg()) cst := mustLoc(t, "America/Chicago") wall := time.Date(2026, 8, 28, 19, 0, 0, 0, cst) // Friday evening gate.SetClock(func() time.Time { return wall }) gate.mu.Lock() gate.lastSnap = &QuotaSnapshot{Account: "zai-1", Source: "usage_url", FetchedAt: wall, Buckets: []Bucket{ {ID: Bucket5h, Used: 12000, Limit: 28000}, {ID: BucketWeekly, Used: 135800, Limit: 140000}, // 97% }} gate.mu.Unlock() for _, class := range []string{"primary", "code", "study", "read", "review"} { d := gate.Decide(context.Background(), class) if d.Action != ActionDefer { t.Fatalf("class %s at 97%% weekly: action = %s, want defer", class, d.Action) } if !strings.Contains(d.Reason, "weekly") || !strings.Contains(d.Reason, "97%") { t.Errorf("defer reason must surface the bucket + ratio: %q", d.Reason) } } } func TestDecideLevels(t *testing.T) { cfg := gateTestCfg() gate := newGate(t, cfg) // Wednesday noon CST: off-peak, mid-week. noon := time.Date(2026, 9, 2, 12, 0, 0, 0, mustLoc(t, "America/Chicago")) gate.SetClock(func() time.Time { return noon }) snap := func(pct5, pctW float64) { gate.mu.Lock() defer gate.mu.Unlock() gate.lastSnap = &QuotaSnapshot{Account: cfg.Account, Source: "usage_url", FetchedAt: noon, Buckets: []Bucket{ {ID: Bucket5h, Used: 28000 * pct5 / 100, Limit: 28000}, {ID: BucketWeekly, Used: 140000 * pctW / 100, Limit: 140000}, }} } cases := []struct { name string pct5, pctW float64 class string want string wantIn string }{ {"healthy allows flagship", 10, 10, "primary", ActionAllow, ""}, {"healthy allows flash", 10, 10, "study", ActionAllow, ""}, {"defer level: heavy defers", 90, 50, "code", ActionDefer, ">= defer"}, {"defer level: flash continues", 90, 50, "study", ActionAllow, ""}, {"block level: flash defers too", 50, 96, "study", ActionDefer, ">= block"}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { snap(c.pct5, c.pctW) d := gate.Decide(context.Background(), c.class) if d.Action != c.want { t.Fatalf("action = %s (%s), want %s", d.Action, d.Reason, c.want) } if c.wantIn != "" && !strings.Contains(d.Reason, c.wantIn) { t.Errorf("reason %q missing %q", d.Reason, c.wantIn) } }) } // Peak window restriction (quota healthy): Friday 03:00 CST. peak := time.Date(2026, 8, 28, 3, 0, 0, 0, mustLoc(t, "America/Chicago")) gate.SetClock(func() time.Time { return peak }) snap(10, 10) if d := gate.Decide(context.Background(), "primary"); d.Action != ActionDefer || !strings.Contains(d.Reason, "peak window") { t.Errorf("peak + flagship class: %+v, want defer/peak reason", d) } if d := gate.Decide(context.Background(), "study"); d.Action != ActionAllow { t.Errorf("peak + flash class must allow: %+v", d) } // Block beats peak: even flash defers when the wall is hit in-peak. snap(10, 97) if d := gate.Decide(context.Background(), "study"); d.Action != ActionDefer { t.Errorf("peak + wall: %+v, want defer", d) } } func gateTestCfg() config.QuotaConfig { return config.QuotaConfig{ Enabled: true, Account: "zai-1", Plan5hCredits: 28000, PlanWeeklyCredits: 140000, PollIntervalSecs: 300, DeferAtPct: 85, BlockAtPct: 95, PeakStart: "01:00", PeakEnd: "05:00", Timezone: "America/Chicago", PeakWeekdaysOnly: true, PeakClasses: []string{"study", "read"}, } } func newGate(t *testing.T, cfg config.QuotaConfig) *Gate { t.Helper() gate, err := NewGate(cfg, nil, nil) if err != nil { t.Fatal(err) } return gate } // (end of quota gate tests)