1package modes
2
3import (
4 "context"
5 "errors"
6 "sync/atomic"
7 "testing"
8 "time"
9
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12)
13
14func newCodexUsageTestInteractive(fetch func(context.Context) (*provider.CodexWeeklyUsage, error)) *Interactive {
15 ag := core.NewAgent(nil, "test", "", nil)
16 BindCodexWeeklyUsageFetcher(ag, fetch)
17 return NewInteractive(InteractiveConfig{Provider: "openai-codex", AuthMethod: "oauth", Agent: ag})
18}
19
20// Wait for the worker's channel, then make that result available to the normal
21// tick path. No sleeps or real five-minute timer are needed.
22func finishCodexUsageRequest(t *testing.T, i *Interactive, now time.Time) {
23 t.Helper()
24 var usage *provider.CodexWeeklyUsage
25 select {
26 case usage = <-i.codexUsage.result:
27 case <-time.After(time.Second):
28 t.Fatal("usage request did not finish")
29 }
30 ready := make(chan *provider.CodexWeeklyUsage, 1)
31 ready <- usage
32 i.codexUsage.result = ready
33 i.refreshCodexUsage(context.Background(), now)
34}
35
36func TestCodexUsagePollIntervalAndFailure(t *testing.T) {
37 now := time.Unix(2000000000, 0)
38 var calls atomic.Int32
39 i := newCodexUsageTestInteractive(func(ctx context.Context) (*provider.CodexWeeklyUsage, error) {
40 if _, ok := ctx.Deadline(); !ok {
41 t.Error("usage request has no deadline")
42 }
43 if calls.Add(1) == 2 {
44 return nil, errors.New("synthetic failure")
45 }
46 return &provider.CodexWeeklyUsage{RemainingPercent: 69}, nil
47 })
48 defer i.resetCodexUsage()
49 i.refreshCodexUsage(context.Background(), now)
50 finishCodexUsageRequest(t, i, now)
51 if got := i.codexWeeklyLabelLocked(now); got != "weekly:69%" {
52 t.Fatalf("label = %q", got)
53 }
54 i.refreshCodexUsage(context.Background(), now.Add(5*time.Minute-time.Nanosecond))
55 if calls.Load() != 1 {
56 t.Fatal("polled before five minutes")
57 }
58 now = now.Add(5 * time.Minute)
59 i.refreshCodexUsage(context.Background(), now)
60 if got := i.codexWeeklyLabelLocked(now); got != "weekly:?" {
61 t.Fatalf("stale usage during refresh: %q", got)
62 }
63 finishCodexUsageRequest(t, i, now)
64 if calls.Load() != 2 || i.codexWeeklyLabelLocked(now) != "weekly:?" || i.statusErr != "" {
65 t.Fatal("failed request should clear only the weekly meter")
66 }
67 now = now.Add(5 * time.Minute)
68 i.refreshCodexUsage(context.Background(), now)
69 finishCodexUsageRequest(t, i, now)
70 if calls.Load() != 3 || i.codexWeeklyLabelLocked(now) != "weekly:69%" {
71 t.Fatal("polling did not recover")
72 }
73}
74
75func TestCodexUsageAsyncSingleFlightAndCancellation(t *testing.T) {
76 started := make(chan context.Context, 2)
77 i := newCodexUsageTestInteractive(func(ctx context.Context) (*provider.CodexWeeklyUsage, error) {
78 started <- ctx
79 <-ctx.Done()
80 return &provider.CodexWeeklyUsage{RemainingPercent: 99}, nil
81 })
82 defer i.resetCodexUsage()
83 now := time.Now()
84 i.refreshCodexUsage(context.Background(), now) // must return while fetch is blocked
85 var requestCtx context.Context
86 select {
87 case requestCtx = <-started:
88 case <-time.After(time.Second):
89 t.Fatal("request did not start")
90 }
91 oldResult := i.codexUsage.result
92 i.refreshCodexUsage(context.Background(), now.Add(10*time.Minute))
93 if i.codexUsage.result != oldResult {
94 t.Fatal("started overlapping request")
95 }
96 candidate := core.NewAgent(nil, "replacement", "", nil)
97 BindCodexWeeklyUsageFetcher(candidate, i.codexUsageFetcherLocked())
98 if requestCtx.Err() != nil || i.codexUsage.result != oldResult {
99 t.Fatal("preparing candidate disturbed live request")
100 }
101 i.ApplySessionAgentWithCompactHandoff(candidate, "openai-codex", "replacement", nil)
102 select {
103 case <-requestCtx.Done():
104 case <-time.After(time.Second):
105 t.Fatal("request not canceled")
106 }
107 select {
108 case <-oldResult:
109 case <-time.After(time.Second):
110 t.Fatal("canceled worker did not exit")
111 }
112 if i.codexWeeklyLabelLocked(now) != "weekly:?" {
113 t.Fatal("old account result was published")
114 }
115}
116
117func TestCodexUsageEligibilityAndProviderSwitch(t *testing.T) {
118 var calls atomic.Int32
119 i := newCodexUsageTestInteractive(func(context.Context) (*provider.CodexWeeklyUsage, error) {
120 calls.Add(1)
121 return &provider.CodexWeeklyUsage{RemainingPercent: 0}, nil
122 })
123 defer i.resetCodexUsage()
124 now := time.Now()
125 for _, p := range []string{"openai", "anthropic", "openai-responses"} {
126 i.cfg.Provider = p
127 i.refreshCodexUsage(context.Background(), now)
128 if i.codexWeeklyLabelLocked(now) != "" || i.codexUsage.result != nil {
129 t.Fatalf("polled for %s", p)
130 }
131 }
132 i.cfg.Provider = "openai-codex"
133 subscriptionAgent := i.agent
134 i.ApplySessionAgent(core.NewAgent(nil, "custom", "", nil), "openai-codex", "custom")
135 i.refreshCodexUsage(context.Background(), now)
136 if calls.Load() != 0 || i.codexWeeklyLabelLocked(now) != "" {
137 t.Fatal("polled disabled route")
138 }
139 i.ApplySessionAgent(subscriptionAgent, "openai-codex", "test")
140 i.cfg.AuthMethod = "apikey" // startup metadata from the previous provider
141 i.refreshCodexUsage(context.Background(), now)
142 finishCodexUsageRequest(t, i, now)
143 if i.codexWeeklyLabelLocked(now) != "weekly:0%" {
144 t.Fatal("switching into Codex did not fetch subscription usage")
145 }
146 i.agent = nil // logout
147 i.refreshCodexUsage(context.Background(), now)
148 if i.codexWeeklyLabelLocked(now) != "" || i.codexUsage.usage != nil {
149 t.Fatal("logout retained usage")
150 }
151}
152
153func TestCodexUsageCandidateIsPrivateUntilSessionCommit(t *testing.T) {
154 now := time.Now()
155 i := newCodexUsageTestInteractive(func(context.Context) (*provider.CodexWeeklyUsage, error) {
156 return &provider.CodexWeeklyUsage{RemainingPercent: 69}, nil
157 })
158 defer i.resetCodexUsage()
159 i.refreshCodexUsage(context.Background(), now)
160 finishCodexUsageRequest(t, i, now)
161 previous := i.codexUsage
162 var candidateCalls atomic.Int32
163 candidate := core.NewAgent(nil, "replacement", "", nil)
164 BindCodexWeeklyUsageFetcher(candidate, func(context.Context) (*provider.CodexWeeklyUsage, error) {
165 candidateCalls.Add(1)
166 return &provider.CodexWeeklyUsage{RemainingPercent: 42}, nil
167 })
168 // A prepared candidate may be discarded by a canceled session load.
169 i.refreshCodexUsage(context.Background(), now.Add(time.Minute))
170 if i.codexUsage.usage != previous.usage || i.codexUsage.next != previous.next ||
171 i.codexWeeklyLabelLocked(now) != "weekly:69%" || candidateCalls.Load() != 0 {
172 t.Fatal("private candidate changed the live meter or polling schedule")
173 }
174 i.ApplySessionAgentWithCompactHandoff(candidate, "openai-codex", "replacement", nil)
175 if i.codexWeeklyLabelLocked(now) != "weekly:?" || !i.codexUsage.next.IsZero() {
176 t.Fatal("commit retained the previous account's sample or schedule")
177 }
178 i.refreshCodexUsage(context.Background(), now)
179 finishCodexUsageRequest(t, i, now)
180 if i.codexWeeklyLabelLocked(now) != "weekly:42%" || candidateCalls.Load() != 1 {
181 t.Fatal("commit did not activate the candidate meter")
182 }
183}
184
185func TestCodexUsageMissingAndExpiredWindow(t *testing.T) {
186 now := time.Now()
187 i := newCodexUsageTestInteractive(func(context.Context) (*provider.CodexWeeklyUsage, error) { return nil, nil })
188 defer i.resetCodexUsage()
189 i.refreshCodexUsage(context.Background(), now)
190 finishCodexUsageRequest(t, i, now)
191 if i.codexWeeklyLabelLocked(now) != "weekly:?" {
192 t.Fatal("missing weekly window presented as zero")
193 }
194 i.codexUsage.usage = &provider.CodexWeeklyUsage{RemainingPercent: 69, ResetsAt: now.Add(time.Minute)}
195 if i.codexWeeklyLabelLocked(now.Add(time.Minute)) != "weekly:?" {
196 t.Fatal("expired window still visible")
197 }
198 i.codexUsage.usage.ResetsAt = time.Time{}
199 if i.codexWeeklyLabelLocked(now.Add(codexUsageInterval)) != "weekly:?" {
200 t.Fatal("stale sample still visible")
201 }
202}