diff options
| author | Peter Stone <thepeterstone@gmail.com> | 2026-06-03 23:01:39 +0000 |
|---|---|---|
| committer | Peter Stone <thepeterstone@gmail.com> | 2026-06-03 23:01:39 +0000 |
| commit | 68cf1e24dd2ca2612b72babc4b35280cd3512f92 (patch) | |
| tree | 14b73f21570a2f42ae729ca7e9676de6628d6819 /internal/budget/budget_test.go | |
| parent | b28cfc6ff288d083f6c8e9c055b69bfcadbceccc (diff) | |
| parent | 7388337d3be5cc7f65ef547e30b2e39884dd165b (diff) | |
merge: integrate oss branch — budget gating, MCP back-channel, events, loopback-only bind
Key features from OSS branch:
- Budget gating: rolling per-provider spend caps with BUDGET_EXCEEDED task state
- Agent MCP back-channel: runners mint tokens; /mcp endpoint resolves them
- Chatbot MCP server at /chatbot/mcp (requires api_token in config)
- Event timeline: unified event stream replacing ad-hoc question/summary flows
- Loopback-only default bind (127.0.0.1:8484); external_bind_allowed=true to expose
- repository_url required on task creation (enforces traceability)
- Auto-checker: spawns verification task after each execution
Conflict resolutions:
- serve.go: keep cfg.Runners.XEnabled() conditional registration from main + agentRegistry from oss
- config.go: keep RunnersConfig from main + BudgetConfig/ExternalBindAllowed from oss
- server.go: handleStaticFiles (base-path rewrite) kept; deduplicated duplicate from both sides
- executor/claude.go: accepted oss deletion (ClaudeRunner replaced by ContainerRunner)
- container_test.go: added noopChannel + AgentChannel parameter to Run call
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Diffstat (limited to 'internal/budget/budget_test.go')
| -rw-r--r-- | internal/budget/budget_test.go | 112 |
1 files changed, 112 insertions, 0 deletions
diff --git a/internal/budget/budget_test.go b/internal/budget/budget_test.go new file mode 100644 index 0000000..c257c91 --- /dev/null +++ b/internal/budget/budget_test.go @@ -0,0 +1,112 @@ +package budget + +import ( + "testing" + "time" +) + +type fakeSpend struct { + spends map[string]float64 + since time.Time +} + +func (f *fakeSpend) SpendByProviderSince(since time.Time) (map[string]float64, error) { + f.since = since + return f.spends, nil +} + +func newAcct(spends map[string]float64, lim Limits) (*Accountant, *fakeSpend) { + src := &fakeSpend{spends: spends} + a := New(src, lim) + a.now = func() time.Time { return time.Date(2026, 5, 26, 12, 0, 0, 0, time.UTC) } + return a, src +} + +func TestNew_DefaultsWindow(t *testing.T) { + a, _ := newAcct(nil, Limits{}) + if a.lim.Window != DefaultWindow { + t.Errorf("want default window %v, got %v", DefaultWindow, a.lim.Window) + } +} + +func TestHeadroom_UnlimitedWhenNoLimit(t *testing.T) { + a, _ := newAcct(map[string]float64{"local": 999}, Limits{PerProvider: map[string]float64{"claude": 10}}) + h, err := a.Headroom("local") + if err != nil { + t.Fatal(err) + } + if h.Limited { + t.Errorf("local should be unlimited, got %+v", h) + } +} + +func TestHeadroom_ComputesRemainingAndFraction(t *testing.T) { + a, src := newAcct(map[string]float64{"claude": 4}, Limits{Window: 5 * time.Hour, PerProvider: map[string]float64{"claude": 10}}) + h, err := a.Headroom("claude") + if err != nil { + t.Fatal(err) + } + if !h.Limited || h.Limit != 10 || h.Spent != 4 || h.Remaining != 6 { + t.Errorf("unexpected headroom: %+v", h) + } + if h.Fraction != 0.6 { + t.Errorf("fraction: want 0.6, got %v", h.Fraction) + } + // Window lookback is now-5h. + want := time.Date(2026, 5, 26, 7, 0, 0, 0, time.UTC) + if !src.since.Equal(want) { + t.Errorf("since: want %v, got %v", want, src.since) + } +} + +func TestHeadroom_ClampsNegativeRemaining(t *testing.T) { + a, _ := newAcct(map[string]float64{"claude": 15}, Limits{PerProvider: map[string]float64{"claude": 10}}) + h, _ := a.Headroom("claude") + if h.Remaining != 0 { + t.Errorf("remaining should clamp to 0, got %v", h.Remaining) + } +} + +func TestAllow(t *testing.T) { + a, _ := newAcct(map[string]float64{"claude": 8}, Limits{PerProvider: map[string]float64{"claude": 10}}) + + if ok, _ := a.Allow("claude", 1.5); !ok { + t.Error("8 + 1.5 <= 10 should be allowed") + } + if ok, _ := a.Allow("claude", 3); ok { + t.Error("8 + 3 > 10 should be denied") + } + if ok, _ := a.Allow("local", 1000); !ok { + t.Error("unlimited provider should always allow") + } +} + +func TestAll_SortedAndSingleQuery(t *testing.T) { + a, src := newAcct(map[string]float64{"claude": 2, "gemini": 1}, + Limits{PerProvider: map[string]float64{"gemini": 5, "claude": 10}}) + calls := 0 + orig := src.spends + a2 := New(&countingSpend{spends: orig, calls: &calls}, a.lim) + a2.now = a.now + + hs, err := a2.All() + if err != nil { + t.Fatal(err) + } + if len(hs) != 2 || hs[0].Provider != "claude" || hs[1].Provider != "gemini" { + t.Errorf("want [claude gemini] sorted, got %+v", hs) + } + if calls != 1 { + t.Errorf("All should issue exactly one spend query, got %d", calls) + } +} + +type countingSpend struct { + spends map[string]float64 + calls *int +} + +func (c *countingSpend) SpendByProviderSince(time.Time) (map[string]float64, error) { + *c.calls++ + return c.spends, nil +} |
