|
| 1 | +package admin |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "encoding/json" |
| 6 | + "errors" |
| 7 | + "net/http" |
| 8 | + "net/http/httptest" |
| 9 | + "strings" |
| 10 | + "testing" |
| 11 | + |
| 12 | + "github.com/Wei-Shaw/sub2api/internal/service" |
| 13 | + "github.com/gin-gonic/gin" |
| 14 | + "github.com/stretchr/testify/require" |
| 15 | +) |
| 16 | + |
| 17 | +const ( |
| 18 | + accountListRoute = "/api/v1/admin/accounts" |
| 19 | + upstreamBillingRatesRoute = "/api/v1/admin/accounts/upstream-billing-rates" |
| 20 | +) |
| 21 | + |
| 22 | +// accountPageCall is one list call as the admin service received it: the scalar arguments of ListAccounts, or the |
| 23 | +// filters of ListAccountsConsole when console is set. |
| 24 | +type accountPageCall struct { |
| 25 | + page, pageSize int |
| 26 | + platform, accountType, status, search string |
| 27 | + groupID int64 |
| 28 | + privacyMode, sortBy, sortOrder string |
| 29 | + console *service.AccountConsoleFilters |
| 30 | +} |
| 31 | + |
| 32 | +// group is the group the call filtered by, whichever list method received it. |
| 33 | +func (c accountPageCall) group() int64 { |
| 34 | + if c.console != nil { |
| 35 | + return c.console.GroupID |
| 36 | + } |
| 37 | + return c.groupID |
| 38 | +} |
| 39 | + |
| 40 | +// billingRatesAdmin records the account page each request resolves and answers every list call with rows and total. |
| 41 | +type billingRatesAdmin struct { |
| 42 | + *accountTaxonomyHandlerStub |
| 43 | + calls []accountPageCall |
| 44 | + rows []service.Account |
| 45 | + total int64 |
| 46 | + listErr error |
| 47 | +} |
| 48 | + |
| 49 | +func (s *billingRatesAdmin) ListAccounts(_ context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string, sortBy, sortOrder string) ([]service.Account, int64, error) { |
| 50 | + s.calls = append(s.calls, accountPageCall{ |
| 51 | + page: page, pageSize: pageSize, platform: platform, accountType: accountType, status: status, search: search, |
| 52 | + groupID: groupID, privacyMode: privacyMode, sortBy: sortBy, sortOrder: sortOrder, |
| 53 | + }) |
| 54 | + return s.rows, s.total, s.listErr |
| 55 | +} |
| 56 | + |
| 57 | +func (s *billingRatesAdmin) ListAccountsConsole(_ context.Context, page, pageSize int, filters service.AccountConsoleFilters) ([]service.Account, int64, error) { |
| 58 | + s.calls = append(s.calls, accountPageCall{page: page, pageSize: pageSize, console: &filters}) |
| 59 | + return s.rows, s.total, s.listErr |
| 60 | +} |
| 61 | + |
| 62 | +func setupUpstreamBillingRatesRouter() (*gin.Engine, *billingRatesAdmin) { |
| 63 | + gin.SetMode(gin.TestMode) |
| 64 | + router := gin.New() |
| 65 | + admin := &billingRatesAdmin{accountTaxonomyHandlerStub: newAccountTaxonomyHandlerStub()} |
| 66 | + handler := NewAccountHandler(admin, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) |
| 67 | + router.GET(accountListRoute, handler.List) |
| 68 | + router.GET(upstreamBillingRatesRoute, handler.GetUpstreamBillingRates) |
| 69 | + return router, admin |
| 70 | +} |
| 71 | + |
| 72 | +// requestAccountPage sends the query to the route and returns the response with the list calls the request made. |
| 73 | +func requestAccountPage(router *gin.Engine, admin *billingRatesAdmin, route, query string) (*httptest.ResponseRecorder, []accountPageCall) { |
| 74 | + admin.calls = nil |
| 75 | + recorder := httptest.NewRecorder() |
| 76 | + router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, route+"?"+query, nil)) |
| 77 | + return recorder, admin.calls |
| 78 | +} |
| 79 | + |
| 80 | +func TestAccountHandlerGetUpstreamBillingRatesResolvesSamePageAsList(t *testing.T) { |
| 81 | + router, admin := setupUpstreamBillingRatesRouter() |
| 82 | + digestTerm := "sha256:" + strings.Repeat("ab", 32) |
| 83 | + over100Bytes := strings.Repeat("k", 120) |
| 84 | + for _, test := range []struct { |
| 85 | + name string |
| 86 | + query string |
| 87 | + console bool |
| 88 | + }{ |
| 89 | + {name: "no_filter"}, |
| 90 | + {name: "search", query: "search=" + digestTerm}, |
| 91 | + {name: "search_over_100_bytes", query: "search=" + over100Bytes}, |
| 92 | + {name: "group", query: "group_id=12"}, |
| 93 | + {name: "ungrouped", query: "group_id=ungrouped"}, |
| 94 | + {name: "privacy_and_billing_rate_sort", query: "privacy_mode=training_off&sort_by=upstream_billing_rate&sort_order=desc"}, |
| 95 | + {name: "scalar_filters", query: "platform=openai&type=apikey&status=active"}, |
| 96 | + {name: "page", query: "page=3&page_size=50"}, |
| 97 | + {name: "statuses", query: "statuses=active,error", console: true}, |
| 98 | + {name: "console_ungrouped", query: "plans=plus&group_id=ungrouped", console: true}, |
| 99 | + {name: "console_search_over_100_bytes", query: "tags=3&search=" + over100Bytes, console: true}, |
| 100 | + {name: "console_page", query: "page=2&page_size=5&account_ids=11,12", console: true}, |
| 101 | + { |
| 102 | + name: "full_console", |
| 103 | + query: "platforms=openai,anthropic&types=apikey,oauth&statuses=active,error&plans=plus,team" + |
| 104 | + "&proxies=direct,4&folder=uncategorized,7&tags=3,9&account_ids=11,12,13&group_id=12" + |
| 105 | + "&privacy_mode=training_off&search=" + digestTerm + "&sort_by=upstream_billing_rate&sort_order=desc", |
| 106 | + console: true, |
| 107 | + }, |
| 108 | + } { |
| 109 | + t.Run(test.name, func(t *testing.T) { |
| 110 | + listResponse, listed := requestAccountPage(router, admin, accountListRoute, test.query) |
| 111 | + ratesResponse, rated := requestAccountPage(router, admin, upstreamBillingRatesRoute, test.query) |
| 112 | + |
| 113 | + require.Equal(t, http.StatusOK, listResponse.Code) |
| 114 | + require.Equal(t, http.StatusOK, ratesResponse.Code) |
| 115 | + require.Len(t, listed, 1) |
| 116 | + require.Equal(t, test.console, listed[0].console != nil, "the query must take the list path the case is about") |
| 117 | + require.Equal(t, listed, rated) |
| 118 | + }) |
| 119 | + } |
| 120 | +} |
| 121 | + |
| 122 | +func TestAccountHandlerGetUpstreamBillingRatesFailsLikeList(t *testing.T) { |
| 123 | + listErr := errors.New("synthetic list outage") |
| 124 | + for _, test := range []struct { |
| 125 | + name string |
| 126 | + query string |
| 127 | + listErr error |
| 128 | + wantStatus int |
| 129 | + wantCalls int |
| 130 | + }{ |
| 131 | + {name: "invalid_group", query: "group_id=abc", wantStatus: http.StatusBadRequest}, |
| 132 | + {name: "negative_group", query: "group_id=-3", wantStatus: http.StatusBadRequest}, |
| 133 | + {name: "console_invalid_group", query: "statuses=active&group_id=abc", wantStatus: http.StatusBadRequest}, |
| 134 | + {name: "console_invalid_filter", query: "tags=abc", wantStatus: http.StatusBadRequest}, |
| 135 | + {name: "list_error", listErr: listErr, wantStatus: http.StatusInternalServerError, wantCalls: 1}, |
| 136 | + {name: "console_list_error", query: "statuses=active", listErr: listErr, wantStatus: http.StatusInternalServerError, wantCalls: 1}, |
| 137 | + } { |
| 138 | + t.Run(test.name, func(t *testing.T) { |
| 139 | + router, admin := setupUpstreamBillingRatesRouter() |
| 140 | + admin.listErr = test.listErr |
| 141 | + listResponse, _ := requestAccountPage(router, admin, accountListRoute, test.query) |
| 142 | + ratesResponse, rated := requestAccountPage(router, admin, upstreamBillingRatesRoute, test.query) |
| 143 | + |
| 144 | + require.Equal(t, test.wantStatus, listResponse.Code) |
| 145 | + require.Equal(t, listResponse.Code, ratesResponse.Code) |
| 146 | + require.JSONEq(t, listResponse.Body.String(), ratesResponse.Body.String()) |
| 147 | + require.Len(t, rated, test.wantCalls) |
| 148 | + }) |
| 149 | + } |
| 150 | +} |
| 151 | + |
| 152 | +func TestAccountHandlerGetUpstreamBillingRatesFallsBackToOlderGroupName(t *testing.T) { |
| 153 | + router, admin := setupUpstreamBillingRatesRouter() |
| 154 | + for _, test := range []struct { |
| 155 | + query string |
| 156 | + wantGroup int64 |
| 157 | + }{ |
| 158 | + {query: "group=12", wantGroup: 12}, |
| 159 | + {query: "group=12&statuses=active", wantGroup: 12}, |
| 160 | + {query: "group=12&group_id=73", wantGroup: 73}, |
| 161 | + {query: "group=12&group_id=73&statuses=active", wantGroup: 73}, |
| 162 | + // An empty group_id is still the name List reads, and means no group filter there. |
| 163 | + {query: "group=12&group_id=", wantGroup: 0}, |
| 164 | + } { |
| 165 | + t.Run(test.query, func(t *testing.T) { |
| 166 | + recorder, calls := requestAccountPage(router, admin, upstreamBillingRatesRoute, test.query) |
| 167 | + |
| 168 | + require.Equal(t, http.StatusOK, recorder.Code) |
| 169 | + require.Len(t, calls, 1) |
| 170 | + require.Equal(t, test.wantGroup, calls[0].group()) |
| 171 | + }) |
| 172 | + } |
| 173 | +} |
| 174 | + |
| 175 | +func TestAccountHandlerGetUpstreamBillingRatesKeepsResponseShapeWithConsoleFilters(t *testing.T) { |
| 176 | + router, admin := setupUpstreamBillingRatesRouter() |
| 177 | + admin.rows = []service.Account{ |
| 178 | + {ID: 42, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, Extra: map[string]any{ |
| 179 | + service.UpstreamBillingProbeExtraKey: map[string]any{ |
| 180 | + "status": service.UpstreamBillingProbeStatusOK, "data": map[string]any{"effective_rate_multiplier": 0.045}, |
| 181 | + "last_attempt_at": "2026-08-01T00:00:00Z", "next_probe_at": "2026-08-01T00:30:00Z", |
| 182 | + }, |
| 183 | + }}, |
| 184 | + {ID: 7, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}, |
| 185 | + {ID: 19, Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth}, |
| 186 | + } |
| 187 | + admin.total = 57 |
| 188 | + const query = "page=2&page_size=3&statuses=active&tags=3&sort_by=upstream_billing_rate&sort_order=desc" |
| 189 | + |
| 190 | + recorder, calls := requestAccountPage(router, admin, upstreamBillingRatesRoute, query) |
| 191 | + require.Equal(t, http.StatusOK, recorder.Code) |
| 192 | + require.Len(t, calls, 1) |
| 193 | + require.NotNil(t, calls[0].console) |
| 194 | + |
| 195 | + var body struct { |
| 196 | + Data json.RawMessage `json:"data"` |
| 197 | + } |
| 198 | + require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body)) |
| 199 | + require.JSONEq(t, `{ |
| 200 | + "items": [ |
| 201 | + {"account_id": 42, "snapshot": { |
| 202 | + "status": "ok", "data": {"effective_rate_multiplier": 0.045}, |
| 203 | + "last_attempt_at": "2026-08-01T00:00:00Z", "next_probe_at": "2026-08-01T00:30:00Z" |
| 204 | + }}, |
| 205 | + {"account_id": 7, "snapshot": null}, |
| 206 | + {"account_id": 19, "snapshot": null} |
| 207 | + ], |
| 208 | + "total": 57, "page": 2, "page_size": 3 |
| 209 | + }`, string(body.Data)) |
| 210 | + require.Equal(t, "private, no-cache", recorder.Header().Get("Cache-Control")) |
| 211 | + etag := recorder.Header().Get("ETag") |
| 212 | + require.NotEmpty(t, etag) |
| 213 | + |
| 214 | + request := httptest.NewRequest(http.MethodGet, upstreamBillingRatesRoute+"?"+query, nil) |
| 215 | + request.Header.Set("If-None-Match", etag) |
| 216 | + notModified := httptest.NewRecorder() |
| 217 | + router.ServeHTTP(notModified, request) |
| 218 | + require.Equal(t, http.StatusNotModified, notModified.Code) |
| 219 | + require.Empty(t, notModified.Body.String()) |
| 220 | +} |
0 commit comments