account_usage_service_batch_test.go 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. package service
  2. import (
  3. "context"
  4. "strings"
  5. "testing"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
  9. )
  10. // Minimal UsageLogRepository stub for batch usage tests (HEAD lacks geminiUsageLogRepoStub).
  11. type usageBatchLogRepoStub struct{}
  12. var _ UsageLogRepository = (*usageBatchLogRepoStub)(nil)
  13. func (r *usageBatchLogRepoStub) Create(context.Context, *UsageLog) (bool, error) {
  14. return false, nil
  15. }
  16. func (r *usageBatchLogRepoStub) GetByID(context.Context, int64) (*UsageLog, error) {
  17. return nil, nil
  18. }
  19. func (r *usageBatchLogRepoStub) Delete(context.Context, int64) error { return nil }
  20. func (r *usageBatchLogRepoStub) ListByUser(context.Context, int64, pagination.PaginationParams) ([]UsageLog, *pagination.PaginationResult, error) {
  21. return nil, nil, nil
  22. }
  23. func (r *usageBatchLogRepoStub) ListByAPIKey(context.Context, int64, pagination.PaginationParams) ([]UsageLog, *pagination.PaginationResult, error) {
  24. return nil, nil, nil
  25. }
  26. func (r *usageBatchLogRepoStub) ListByAccount(context.Context, int64, pagination.PaginationParams) ([]UsageLog, *pagination.PaginationResult, error) {
  27. return nil, nil, nil
  28. }
  29. func (r *usageBatchLogRepoStub) ListByUserAndTimeRange(context.Context, int64, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
  30. return nil, nil, nil
  31. }
  32. func (r *usageBatchLogRepoStub) ListByAPIKeyAndTimeRange(context.Context, int64, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
  33. return nil, nil, nil
  34. }
  35. func (r *usageBatchLogRepoStub) ListByAccountAndTimeRange(context.Context, int64, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
  36. return nil, nil, nil
  37. }
  38. func (r *usageBatchLogRepoStub) ListByModelAndTimeRange(context.Context, string, time.Time, time.Time) ([]UsageLog, *pagination.PaginationResult, error) {
  39. return nil, nil, nil
  40. }
  41. func (r *usageBatchLogRepoStub) GetAccountWindowStats(context.Context, int64, time.Time) (*usagestats.AccountStats, error) {
  42. return &usagestats.AccountStats{}, nil
  43. }
  44. func (r *usageBatchLogRepoStub) GetAccountTodayStats(context.Context, int64) (*usagestats.AccountStats, error) {
  45. return &usagestats.AccountStats{}, nil
  46. }
  47. func (r *usageBatchLogRepoStub) GetDashboardStats(context.Context) (*usagestats.DashboardStats, error) {
  48. return nil, nil
  49. }
  50. func (r *usageBatchLogRepoStub) GetUsageTrendWithFilters(context.Context, time.Time, time.Time, string, int64, int64, int64, int64, string, *int16, *bool, *int8) ([]usagestats.TrendDataPoint, error) {
  51. return nil, nil
  52. }
  53. func (r *usageBatchLogRepoStub) GetModelStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, *int16, *bool, *int8) ([]usagestats.ModelStat, error) {
  54. return nil, nil
  55. }
  56. func (r *usageBatchLogRepoStub) GetEndpointStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, string, *int16, *bool, *int8) ([]usagestats.EndpointStat, error) {
  57. return nil, nil
  58. }
  59. func (r *usageBatchLogRepoStub) GetUpstreamEndpointStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, string, *int16, *bool, *int8) ([]usagestats.EndpointStat, error) {
  60. return nil, nil
  61. }
  62. func (r *usageBatchLogRepoStub) GetGroupStatsWithFilters(context.Context, time.Time, time.Time, int64, int64, int64, int64, *int16, *bool, *int8) ([]usagestats.GroupStat, error) {
  63. return nil, nil
  64. }
  65. func (r *usageBatchLogRepoStub) GetUserBreakdownStats(context.Context, time.Time, time.Time, usagestats.UserBreakdownDimension, int) ([]usagestats.UserBreakdownItem, error) {
  66. return nil, nil
  67. }
  68. func (r *usageBatchLogRepoStub) GetAllGroupUsageSummary(context.Context, time.Time) ([]usagestats.GroupUsageSummary, error) {
  69. return nil, nil
  70. }
  71. func (r *usageBatchLogRepoStub) GetAPIKeyUsageTrend(context.Context, time.Time, time.Time, string, int) ([]usagestats.APIKeyUsageTrendPoint, error) {
  72. return nil, nil
  73. }
  74. func (r *usageBatchLogRepoStub) GetUserUsageTrend(context.Context, time.Time, time.Time, string, int) ([]usagestats.UserUsageTrendPoint, error) {
  75. return nil, nil
  76. }
  77. func (r *usageBatchLogRepoStub) GetUserSpendingRanking(context.Context, time.Time, time.Time, int) (*usagestats.UserSpendingRankingResponse, error) {
  78. return nil, nil
  79. }
  80. func (r *usageBatchLogRepoStub) GetBatchUserUsageStats(context.Context, []int64, time.Time, time.Time) (map[int64]*usagestats.BatchUserUsageStats, error) {
  81. return nil, nil
  82. }
  83. func (r *usageBatchLogRepoStub) GetBatchAPIKeyUsageStats(context.Context, []int64, time.Time, time.Time) (map[int64]*usagestats.BatchAPIKeyUsageStats, error) {
  84. return nil, nil
  85. }
  86. func (r *usageBatchLogRepoStub) GetUserDashboardStats(context.Context, int64) (*usagestats.UserDashboardStats, error) {
  87. return nil, nil
  88. }
  89. func (r *usageBatchLogRepoStub) GetAPIKeyDashboardStats(context.Context, int64) (*usagestats.UserDashboardStats, error) {
  90. return nil, nil
  91. }
  92. func (r *usageBatchLogRepoStub) GetUserUsageTrendByUserID(context.Context, int64, time.Time, time.Time, string) ([]usagestats.TrendDataPoint, error) {
  93. return nil, nil
  94. }
  95. func (r *usageBatchLogRepoStub) GetUserModelStats(context.Context, int64, time.Time, time.Time) ([]usagestats.ModelStat, error) {
  96. return nil, nil
  97. }
  98. func (r *usageBatchLogRepoStub) ListWithFilters(context.Context, pagination.PaginationParams, usagestats.UsageLogFilters) ([]UsageLog, *pagination.PaginationResult, error) {
  99. return nil, nil, nil
  100. }
  101. func (r *usageBatchLogRepoStub) GetGlobalStats(context.Context, time.Time, time.Time) (*usagestats.UsageStats, error) {
  102. return nil, nil
  103. }
  104. func (r *usageBatchLogRepoStub) GetStatsWithFilters(context.Context, usagestats.UsageLogFilters) (*usagestats.UsageStats, error) {
  105. return nil, nil
  106. }
  107. func (r *usageBatchLogRepoStub) GetAccountUsageStats(context.Context, int64, time.Time, time.Time) (*usagestats.AccountUsageStatsResponse, error) {
  108. return nil, nil
  109. }
  110. func (r *usageBatchLogRepoStub) GetUserStatsAggregated(context.Context, int64, time.Time, time.Time) (*usagestats.UsageStats, error) {
  111. return nil, nil
  112. }
  113. func (r *usageBatchLogRepoStub) GetAPIKeyStatsAggregated(context.Context, int64, time.Time, time.Time) (*usagestats.UsageStats, error) {
  114. return nil, nil
  115. }
  116. func (r *usageBatchLogRepoStub) GetAccountStatsAggregated(context.Context, int64, time.Time, time.Time) (*usagestats.UsageStats, error) {
  117. return nil, nil
  118. }
  119. func (r *usageBatchLogRepoStub) GetModelStatsAggregated(context.Context, string, time.Time, time.Time) (*usagestats.UsageStats, error) {
  120. return nil, nil
  121. }
  122. func (r *usageBatchLogRepoStub) GetDailyStatsAggregated(context.Context, int64, time.Time, time.Time) ([]map[string]any, error) {
  123. return nil, nil
  124. }
  125. func TestAccountUsageService_GetUsageBatch_BestEffortByAccount(t *testing.T) {
  126. t.Parallel()
  127. resetAt := time.Now().Add(2 * time.Hour).UTC().Truncate(time.Second)
  128. repo := &stubOpenAIAccountRepo{
  129. accounts: []Account{
  130. {
  131. ID: 7001,
  132. Platform: PlatformAnthropic,
  133. Type: AccountTypeOAuth,
  134. Extra: map[string]any{
  135. "passive_usage_7d_utilization": 0.62,
  136. },
  137. },
  138. {
  139. ID: 7002,
  140. Platform: PlatformOpenAI,
  141. Type: AccountTypeOAuth,
  142. Extra: map[string]any{
  143. "codex_usage_updated_at": time.Now().UTC().Format(time.RFC3339),
  144. "codex_5h_used_percent": 18.0,
  145. "codex_5h_reset_at": resetAt.Format(time.RFC3339),
  146. "codex_7d_used_percent": 34.0,
  147. "codex_7d_reset_at": resetAt.Add(24 * time.Hour).Format(time.RFC3339),
  148. "workspace_id": "org-test",
  149. "chatgpt_account_id": "acct-test",
  150. "openai_snapshot_version": "test",
  151. },
  152. },
  153. {
  154. ID: 7003,
  155. Platform: PlatformOpenAI,
  156. Type: AccountTypeAPIKey,
  157. },
  158. },
  159. }
  160. svc := &AccountUsageService{
  161. accountRepo: repo,
  162. usageLogRepo: &usageBatchLogRepoStub{},
  163. cache: NewUsageCache(),
  164. }
  165. usageByAccount, errorsByAccount, err := svc.GetUsageBatch(context.Background(), []int64{7001, 7002, 7003, 7002}, false)
  166. if err != nil {
  167. t.Fatalf("GetUsageBatch() error = %v", err)
  168. }
  169. if usageByAccount[7001] == nil || usageByAccount[7001].Source != "passive" {
  170. t.Fatalf("expected anthropic passive usage, got %#v", usageByAccount[7001])
  171. }
  172. if usageByAccount[7002] == nil || usageByAccount[7002].FiveHour == nil || usageByAccount[7002].FiveHour.Utilization != 18.0 {
  173. t.Fatalf("expected openai snapshot usage, got %#v", usageByAccount[7002])
  174. }
  175. if !strings.Contains(strings.ToLower(errorsByAccount[7003]), "does not support usage query") {
  176. t.Fatalf("expected API key account error to be preserved, got %q", errorsByAccount[7003])
  177. }
  178. }