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
+}