1package subagents
2
3import (
4 "testing"
5
6 "github.com/bnema/zut/packages/provider"
7)
8
9func TestContextBudgetLimitUsesFullModelContext(t *testing.T) {
10 limit, err := ContextBudgetLimit(500_000)
11 if err != nil {
12 t.Fatal(err)
13 }
14 if limit != 500_000 {
15 t.Fatalf("ContextBudgetLimit = %d, want 500000", limit)
16 }
17 if _, err := ContextBudgetLimit(0); err == nil {
18 t.Fatal("ContextBudgetLimit(0) succeeded")
19 }
20}
21
22func TestValidateResidentBudgetRequiresModelContext(t *testing.T) {
23 if err := ValidateResidentBudget(ResidentChildSpec{BudgetLimit: 500_000, BudgetSource: "model_context"}); err != nil {
24 t.Fatal(err)
25 }
26 for _, spec := range []ResidentChildSpec{
27 {},
28 {BudgetLimit: 500_000},
29 {BudgetLimit: 375_000, BudgetSource: "default_ratio"},
30 } {
31 if err := ValidateResidentBudget(spec); err == nil {
32 t.Fatalf("ValidateResidentBudget(%#v) succeeded", spec)
33 }
34 }
35}
36
37func TestWeightedBudgetUsageDiscountsCacheReads(t *testing.T) {
38 usage := provider.Usage{InputTokens: 100, CacheReadTokens: 100, CacheWriteTokens: 20, OutputTokens: 30, ReasoningTokens: 10, ReasoningTokensKnown: true}
39 if got := WeightedBudgetUsage(usage); got != 175 {
40 t.Fatalf("WeightedBudgetUsage = %d, want 175", got)
41 }
42}
43
44func TestRolloutBudgetOnlyExhaustsAtLimit(t *testing.T) {
45 budget := NewRolloutBudget(1_000, provider.Usage{})
46 for _, used := range []int{500, 700, 850, 900, 999} {
47 budget.Observe(provider.Usage{InputTokens: used})
48 if got := budget.Snapshot().State; got != BudgetNormal {
49 t.Fatalf("usage %d state = %q, want %q", used, got, BudgetNormal)
50 }
51 }
52 budget.Observe(provider.Usage{InputTokens: 1_000})
53 if got := budget.Snapshot().State; got != BudgetExceeded {
54 t.Fatalf("exhausted state = %q, want %q", got, BudgetExceeded)
55 }
56}