1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
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
}
|
