usage_billing_repo_integration_test.go 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "fmt"
  6. "strings"
  7. "testing"
  8. "time"
  9. "github.com/google/uuid"
  10. "github.com/stretchr/testify/require"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. )
  13. func TestUsageBillingRepositoryApply_DeduplicatesBalanceBilling(t *testing.T) {
  14. ctx := context.Background()
  15. client := testEntClient(t)
  16. repo := NewUsageBillingRepository(client, integrationDB)
  17. user := mustCreateUser(t, client, &service.User{
  18. Email: fmt.Sprintf("usage-billing-user-%d@example.com", time.Now().UnixNano()),
  19. PasswordHash: "hash",
  20. Balance: 100,
  21. })
  22. apiKey := mustCreateApiKey(t, client, &service.APIKey{
  23. UserID: user.ID,
  24. Key: "sk-usage-billing-" + uuid.NewString(),
  25. Name: "billing",
  26. Quota: 1,
  27. })
  28. account := mustCreateAccount(t, client, &service.Account{
  29. Name: "usage-billing-account-" + uuid.NewString(),
  30. Type: service.AccountTypeAPIKey,
  31. })
  32. requestID := uuid.NewString()
  33. cmd := &service.UsageBillingCommand{
  34. RequestID: requestID,
  35. APIKeyID: apiKey.ID,
  36. UserID: user.ID,
  37. AccountID: account.ID,
  38. AccountType: service.AccountTypeAPIKey,
  39. BalanceCost: 1.25,
  40. APIKeyQuotaCost: 1.25,
  41. APIKeyRateLimitCost: 1.25,
  42. }
  43. result1, err := repo.Apply(ctx, cmd)
  44. require.NoError(t, err)
  45. require.NotNil(t, result1)
  46. require.True(t, result1.Applied)
  47. require.True(t, result1.APIKeyQuotaExhausted)
  48. result2, err := repo.Apply(ctx, cmd)
  49. require.NoError(t, err)
  50. require.NotNil(t, result2)
  51. require.False(t, result2.Applied)
  52. var balance float64
  53. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
  54. require.InDelta(t, 98.75, balance, 0.000001)
  55. var quotaUsed float64
  56. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT quota_used FROM api_keys WHERE id = $1", apiKey.ID).Scan(&quotaUsed))
  57. require.InDelta(t, 1.25, quotaUsed, 0.000001)
  58. var usage5h float64
  59. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT usage_5h FROM api_keys WHERE id = $1", apiKey.ID).Scan(&usage5h))
  60. require.InDelta(t, 1.25, usage5h, 0.000001)
  61. var status string
  62. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT status FROM api_keys WHERE id = $1", apiKey.ID).Scan(&status))
  63. require.Equal(t, service.StatusAPIKeyQuotaExhausted, status)
  64. var dedupCount int
  65. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount))
  66. require.Equal(t, 1, dedupCount)
  67. }
  68. func TestUsageBillingRepositoryApply_DeduplicatesSubscriptionBilling(t *testing.T) {
  69. ctx := context.Background()
  70. client := testEntClient(t)
  71. repo := NewUsageBillingRepository(client, integrationDB)
  72. user := mustCreateUser(t, client, &service.User{
  73. Email: fmt.Sprintf("usage-billing-sub-user-%d@example.com", time.Now().UnixNano()),
  74. PasswordHash: "hash",
  75. })
  76. group := mustCreateGroup(t, client, &service.Group{
  77. Name: "usage-billing-group-" + uuid.NewString(),
  78. Platform: service.PlatformAnthropic,
  79. SubscriptionType: service.SubscriptionTypeSubscription,
  80. })
  81. apiKey := mustCreateApiKey(t, client, &service.APIKey{
  82. UserID: user.ID,
  83. GroupID: &group.ID,
  84. Key: "sk-usage-billing-sub-" + uuid.NewString(),
  85. Name: "billing-sub",
  86. })
  87. subscription := mustCreateSubscription(t, client, &service.UserSubscription{
  88. UserID: user.ID,
  89. GroupID: group.ID,
  90. })
  91. requestID := uuid.NewString()
  92. cmd := &service.UsageBillingCommand{
  93. RequestID: requestID,
  94. APIKeyID: apiKey.ID,
  95. UserID: user.ID,
  96. AccountID: 0,
  97. SubscriptionID: &subscription.ID,
  98. SubscriptionCost: 2.5,
  99. }
  100. result1, err := repo.Apply(ctx, cmd)
  101. require.NoError(t, err)
  102. require.True(t, result1.Applied)
  103. result2, err := repo.Apply(ctx, cmd)
  104. require.NoError(t, err)
  105. require.False(t, result2.Applied)
  106. var dailyUsage float64
  107. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT daily_usage_usd FROM user_subscriptions WHERE id = $1", subscription.ID).Scan(&dailyUsage))
  108. require.InDelta(t, 2.5, dailyUsage, 0.000001)
  109. }
  110. func TestUsageBillingRepositoryApply_RequestFingerprintConflict(t *testing.T) {
  111. ctx := context.Background()
  112. client := testEntClient(t)
  113. repo := NewUsageBillingRepository(client, integrationDB)
  114. user := mustCreateUser(t, client, &service.User{
  115. Email: fmt.Sprintf("usage-billing-conflict-user-%d@example.com", time.Now().UnixNano()),
  116. PasswordHash: "hash",
  117. Balance: 100,
  118. })
  119. apiKey := mustCreateApiKey(t, client, &service.APIKey{
  120. UserID: user.ID,
  121. Key: "sk-usage-billing-conflict-" + uuid.NewString(),
  122. Name: "billing-conflict",
  123. })
  124. requestID := uuid.NewString()
  125. _, err := repo.Apply(ctx, &service.UsageBillingCommand{
  126. RequestID: requestID,
  127. APIKeyID: apiKey.ID,
  128. UserID: user.ID,
  129. BalanceCost: 1.25,
  130. })
  131. require.NoError(t, err)
  132. _, err = repo.Apply(ctx, &service.UsageBillingCommand{
  133. RequestID: requestID,
  134. APIKeyID: apiKey.ID,
  135. UserID: user.ID,
  136. BalanceCost: 2.50,
  137. })
  138. require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict)
  139. }
  140. func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) {
  141. ctx := context.Background()
  142. client := testEntClient(t)
  143. repo := NewUsageBillingRepository(client, integrationDB)
  144. user := mustCreateUser(t, client, &service.User{
  145. Email: fmt.Sprintf("usage-billing-account-user-%d@example.com", time.Now().UnixNano()),
  146. PasswordHash: "hash",
  147. })
  148. apiKey := mustCreateApiKey(t, client, &service.APIKey{
  149. UserID: user.ID,
  150. Key: "sk-usage-billing-account-" + uuid.NewString(),
  151. Name: "billing-account",
  152. })
  153. account := mustCreateAccount(t, client, &service.Account{
  154. Name: "usage-billing-account-quota-" + uuid.NewString(),
  155. Type: service.AccountTypeAPIKey,
  156. Extra: map[string]any{
  157. "quota_limit": 100.0,
  158. },
  159. })
  160. _, err := repo.Apply(ctx, &service.UsageBillingCommand{
  161. RequestID: uuid.NewString(),
  162. APIKeyID: apiKey.ID,
  163. UserID: user.ID,
  164. AccountID: account.ID,
  165. AccountType: service.AccountTypeAPIKey,
  166. AccountQuotaCost: 3.5,
  167. })
  168. require.NoError(t, err)
  169. var quotaUsed float64
  170. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COALESCE((extra->>'quota_used')::numeric, 0) FROM accounts WHERE id = $1", account.ID).Scan(&quotaUsed))
  171. require.InDelta(t, 3.5, quotaUsed, 0.000001)
  172. }
  173. func TestUsageBillingRepositoryApply_EnqueuesSchedulerOutboxOnQuotaCrossing(t *testing.T) {
  174. ctx := context.Background()
  175. client := testEntClient(t)
  176. repo := NewUsageBillingRepository(client, integrationDB)
  177. newFixture := func(t *testing.T, extra map[string]any) (int64, int64) {
  178. t.Helper()
  179. user := mustCreateUser(t, client, &service.User{
  180. Email: fmt.Sprintf("usage-billing-outbox-user-%d-%s@example.com", time.Now().UnixNano(), uuid.NewString()),
  181. PasswordHash: "hash",
  182. })
  183. apiKey := mustCreateApiKey(t, client, &service.APIKey{
  184. UserID: user.ID,
  185. Key: "sk-usage-billing-outbox-" + uuid.NewString(),
  186. Name: "billing-outbox",
  187. })
  188. account := mustCreateAccount(t, client, &service.Account{
  189. Name: "usage-billing-outbox-" + uuid.NewString(),
  190. Type: service.AccountTypeAPIKey,
  191. Extra: extra,
  192. })
  193. return apiKey.ID, account.ID
  194. }
  195. outboxCountFor := func(t *testing.T, accountID int64) int {
  196. t.Helper()
  197. var count int
  198. require.NoError(t, integrationDB.QueryRowContext(ctx,
  199. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  200. service.SchedulerOutboxEventAccountChanged, accountID,
  201. ).Scan(&count))
  202. return count
  203. }
  204. t.Run("daily_first_crossing_enqueues", func(t *testing.T) {
  205. apiKeyID, accountID := newFixture(t, map[string]any{
  206. "quota_daily_limit": 10.0,
  207. })
  208. // 第一次低于日限额:不应入队 outbox
  209. _, err := repo.Apply(ctx, &service.UsageBillingCommand{
  210. RequestID: uuid.NewString(),
  211. APIKeyID: apiKeyID,
  212. AccountID: accountID,
  213. AccountType: service.AccountTypeAPIKey,
  214. AccountQuotaCost: 4,
  215. })
  216. require.NoError(t, err)
  217. require.Equal(t, 0, outboxCountFor(t, accountID), "below limit should not enqueue")
  218. // 第二次跨越日限额:应入队一次 outbox
  219. _, err = repo.Apply(ctx, &service.UsageBillingCommand{
  220. RequestID: uuid.NewString(),
  221. APIKeyID: apiKeyID,
  222. AccountID: accountID,
  223. AccountType: service.AccountTypeAPIKey,
  224. AccountQuotaCost: 8,
  225. })
  226. require.NoError(t, err)
  227. require.Equal(t, 1, outboxCountFor(t, accountID), "crossing daily limit should enqueue once")
  228. // 再次递增(已超):不应重复入队
  229. _, err = repo.Apply(ctx, &service.UsageBillingCommand{
  230. RequestID: uuid.NewString(),
  231. APIKeyID: apiKeyID,
  232. AccountID: accountID,
  233. AccountType: service.AccountTypeAPIKey,
  234. AccountQuotaCost: 2,
  235. })
  236. require.NoError(t, err)
  237. require.Equal(t, 1, outboxCountFor(t, accountID), "subsequent increments beyond limit should not re-enqueue")
  238. })
  239. t.Run("weekly_first_crossing_enqueues", func(t *testing.T) {
  240. apiKeyID, accountID := newFixture(t, map[string]any{
  241. "quota_weekly_limit": 10.0,
  242. })
  243. _, err := repo.Apply(ctx, &service.UsageBillingCommand{
  244. RequestID: uuid.NewString(),
  245. APIKeyID: apiKeyID,
  246. AccountID: accountID,
  247. AccountType: service.AccountTypeAPIKey,
  248. AccountQuotaCost: 15, // 单次即跨越
  249. })
  250. require.NoError(t, err)
  251. require.Equal(t, 1, outboxCountFor(t, accountID), "single-shot crossing weekly limit should enqueue once")
  252. })
  253. }
  254. func TestDashboardAggregationRepositoryCleanupUsageBillingDedup_BatchDeletesOldRows(t *testing.T) {
  255. ctx := context.Background()
  256. repo := newDashboardAggregationRepositoryWithSQL(integrationDB)
  257. oldRequestID := "dedup-old-" + uuid.NewString()
  258. newRequestID := "dedup-new-" + uuid.NewString()
  259. oldCreatedAt := time.Now().UTC().AddDate(0, 0, -400)
  260. newCreatedAt := time.Now().UTC().Add(-time.Hour)
  261. _, err := integrationDB.ExecContext(ctx, `
  262. INSERT INTO usage_billing_dedup (request_id, api_key_id, request_fingerprint, created_at)
  263. VALUES ($1, 1, $2, $3), ($4, 1, $5, $6)
  264. `,
  265. oldRequestID, strings.Repeat("a", 64), oldCreatedAt,
  266. newRequestID, strings.Repeat("b", 64), newCreatedAt,
  267. )
  268. require.NoError(t, err)
  269. require.NoError(t, repo.CleanupUsageBillingDedup(ctx, time.Now().UTC().AddDate(0, 0, -365)))
  270. var oldCount int
  271. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1", oldRequestID).Scan(&oldCount))
  272. require.Equal(t, 0, oldCount)
  273. var newCount int
  274. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1", newRequestID).Scan(&newCount))
  275. require.Equal(t, 1, newCount)
  276. var archivedCount int
  277. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup_archive WHERE request_id = $1", oldRequestID).Scan(&archivedCount))
  278. require.Equal(t, 1, archivedCount)
  279. }
  280. func TestUsageBillingRepositoryApply_DeduplicatesAgainstArchivedKey(t *testing.T) {
  281. ctx := context.Background()
  282. client := testEntClient(t)
  283. repo := NewUsageBillingRepository(client, integrationDB)
  284. aggRepo := newDashboardAggregationRepositoryWithSQL(integrationDB)
  285. user := mustCreateUser(t, client, &service.User{
  286. Email: fmt.Sprintf("usage-billing-archive-user-%d@example.com", time.Now().UnixNano()),
  287. PasswordHash: "hash",
  288. Balance: 100,
  289. })
  290. apiKey := mustCreateApiKey(t, client, &service.APIKey{
  291. UserID: user.ID,
  292. Key: "sk-usage-billing-archive-" + uuid.NewString(),
  293. Name: "billing-archive",
  294. })
  295. requestID := uuid.NewString()
  296. cmd := &service.UsageBillingCommand{
  297. RequestID: requestID,
  298. APIKeyID: apiKey.ID,
  299. UserID: user.ID,
  300. BalanceCost: 1.25,
  301. }
  302. result1, err := repo.Apply(ctx, cmd)
  303. require.NoError(t, err)
  304. require.True(t, result1.Applied)
  305. _, err = integrationDB.ExecContext(ctx, `
  306. UPDATE usage_billing_dedup
  307. SET created_at = $1
  308. WHERE request_id = $2 AND api_key_id = $3
  309. `, time.Now().UTC().AddDate(0, 0, -400), requestID, apiKey.ID)
  310. require.NoError(t, err)
  311. require.NoError(t, aggRepo.CleanupUsageBillingDedup(ctx, time.Now().UTC().AddDate(0, 0, -365)))
  312. result2, err := repo.Apply(ctx, cmd)
  313. require.NoError(t, err)
  314. require.False(t, result2.Applied)
  315. var balance float64
  316. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
  317. require.InDelta(t, 98.75, balance, 0.000001)
  318. }