1package provider
2
3import (
4 "context"
5 "fmt"
6 "net/http"
7 "net/http/httptest"
8 "strings"
9 "sync/atomic"
10 "testing"
11 "time"
12)
13
14func TestFetchCodexWeeklyUsage(t *testing.T) {
15 for _, tc := range []struct {
16 name, body string
17 status int
18 want *float64
19 wantErr bool
20 }{
21 {"secondary weekly", `{"rate_limit":{"primary_window":{"used_percent":12,"limit_window_seconds":18000},"secondary_window":{"used_percent":31,"limit_window_seconds":604800,"reset_at":2000000000}}}`, 200, usagePercent(69), false},
22 {"primary weekly", `{"rate_limit":{"primary_window":{"used_percent":0,"limit_window_seconds":604800}}}`, 200, usagePercent(100), false},
23 {"exhausted", `{"rate_limit":{"secondary_window":{"used_percent":100,"limit_window_seconds":604800}}}`, 200, usagePercent(0), false},
24 {"fractional", `{"rate_limit":{"secondary_window":{"used_percent":30.5,"limit_window_seconds":604800}}}`, 200, usagePercent(69.5), false},
25 {"absent", `{}`, 200, nil, false},
26 {"null", `{"rate_limit":null}`, 200, nil, false},
27 {"non weekly secondary", `{"rate_limit":{"secondary_window":{"used_percent":31,"limit_window_seconds":3600}}}`, 200, nil, false},
28 {"missing percent", `{"rate_limit":{"secondary_window":{"limit_window_seconds":604800}}}`, 200, nil, true},
29 {"negative", `{"rate_limit":{"secondary_window":{"used_percent":-1,"limit_window_seconds":604800}}}`, 200, nil, true},
30 {"over limit", `{"rate_limit":{"secondary_window":{"used_percent":101,"limit_window_seconds":604800}}}`, 200, nil, true},
31 {"invalid json", `private-response`, 200, nil, true},
32 {"oversized", strings.Repeat(" ", 1<<20) + `{}`, 200, nil, true},
33 {"unauthorized", `private-response`, 401, nil, true},
34 {"rate limited", `private-response`, 429, nil, true},
35 {"server error", `private-response`, 503, nil, true},
36 } {
37 t.Run(tc.name, func(t *testing.T) {
38 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
39 if r.Method != http.MethodGet || r.URL.Path != "/backend-api/wham/usage" {
40 t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path)
41 }
42 if r.Header.Get("Authorization") != "Bearer synthetic-token" || r.Header.Get("ChatGPT-Account-Id") != "synthetic-account" || r.Header.Get("User-Agent") != "codex-cli" {
43 t.Error("missing subscription headers")
44 }
45 w.WriteHeader(tc.status)
46 fmt.Fprint(w, tc.body)
47 }))
48 defer srv.Close()
49 got, err := fetchCodexWeeklyUsage(context.Background(), srv.Client(), srv.URL+"/backend-api/wham/usage", "synthetic-token", "synthetic-account")
50 if (err != nil) != tc.wantErr {
51 t.Fatalf("error = %v, wantErr %t", err, tc.wantErr)
52 }
53 if err != nil && (strings.Contains(err.Error(), "synthetic") || strings.Contains(err.Error(), "private-response")) {
54 t.Fatal("error leaked sensitive data")
55 }
56 if tc.want == nil {
57 if got != nil {
58 t.Fatalf("unexpected usage: %+v", got)
59 }
60 return
61 }
62 if got == nil || got.RemainingPercent != *tc.want {
63 t.Fatalf("usage = %+v, want %v", got, *tc.want)
64 }
65 if tc.name == "secondary weekly" && !got.ResetsAt.Equal(time.Unix(2000000000, 0)) {
66 t.Fatalf("reset = %v", got.ResetsAt)
67 }
68 })
69 }
70}
71
72func usagePercent(n float64) *float64 { return &n }
73
74func TestFetchCodexWeeklyUsageCancellation(t *testing.T) {
75 started := make(chan struct{})
76 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
77 close(started)
78 <-r.Context().Done()
79 }))
80 defer srv.Close()
81 ctx, cancel := context.WithCancel(context.Background())
82 defer cancel()
83 done := make(chan error, 1)
84 go func() {
85 _, err := fetchCodexWeeklyUsage(ctx, srv.Client(), srv.URL, "synthetic-token", "synthetic-account")
86 done <- err
87 }()
88 select {
89 case <-started:
90 case <-time.After(time.Second):
91 t.Fatal("request did not start")
92 }
93 cancel()
94 select {
95 case err := <-done:
96 if err == nil {
97 t.Fatal("expected cancellation error")
98 }
99 case <-time.After(time.Second):
100 t.Fatal("request ignored cancellation")
101 }
102}
103
104func TestFetchCodexWeeklyUsageRejectsRedirectsAndMissingCredentials(t *testing.T) {
105 var calls atomic.Int32
106 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
107 calls.Add(1)
108 http.Redirect(w, r, "/private-target", http.StatusFound)
109 }))
110 defer srv.Close()
111 for _, creds := range [][2]string{{"", "account"}, {"token", ""}} {
112 if _, err := fetchCodexWeeklyUsage(context.Background(), srv.Client(), srv.URL, creds[0], creds[1]); err == nil {
113 t.Fatal("missing credentials accepted")
114 }
115 }
116 if calls.Load() != 0 {
117 t.Fatal("missing credentials made a request")
118 }
119 if _, err := fetchCodexWeeklyUsage(context.Background(), srv.Client(), srv.URL, "token", "account"); err == nil {
120 t.Fatal("redirect accepted")
121 }
122 if calls.Load() != 1 {
123 t.Fatalf("redirect followed: %d requests", calls.Load())
124 }
125}