batch_image_settlement_test.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "fmt"
  7. "strings"
  8. "testing"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestBatchImageSettlementService_SettlesAndChargesSuccessfulImagesOnly(t *testing.T) {
  12. repo := newFakeBatchImageRepository()
  13. job := testSettlingBatchImageJob("imgbatch_settle")
  14. job.SuccessCount = 3
  15. job.FailCount = 2
  16. job.ItemCount = 5
  17. job.SessionID = batchImageStringPtr("batch-settlement-session")
  18. repo.jobs[job.BatchID] = job
  19. billing := &fakeBatchImageBillingRepo{}
  20. usageLogs := &openAIRecordUsageLogRepoStub{}
  21. svc := &BatchImageSettlementService{
  22. Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25},
  23. UsageLogRepo: usageLogs,
  24. }
  25. result, err := svc.Settle(context.Background(), job.BatchID)
  26. require.NoError(t, err)
  27. require.Equal(t, 0.75, result.ActualCost)
  28. require.Equal(t, BatchImageCaptureRequestID(job.BatchID), result.RequestID)
  29. require.False(t, result.AlreadySettled)
  30. require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status)
  31. require.NotNil(t, repo.jobs[job.BatchID].ActualCost)
  32. require.Equal(t, 0.75, *repo.jobs[job.BatchID].ActualCost)
  33. require.NotEmpty(t, batchImageDerefString(repo.jobs[job.BatchID].ManifestHash))
  34. require.NotNil(t, repo.jobs[job.BatchID].SettledAt)
  35. require.Equal(t, "batch-settlement-session", batchImageDerefString(usageLogs.lastLog.SessionID))
  36. require.Len(t, billing.captures, 1)
  37. require.Equal(t, int64(321), billing.captures[0].APIKeyID)
  38. require.Equal(t, job.UserID, billing.captures[0].UserID)
  39. require.Equal(t, job.BatchID, billing.captures[0].BatchID)
  40. require.Equal(t, 0.75, billing.captures[0].ActualAmount)
  41. require.Equal(t, 1.25, billing.captures[0].HoldAmount)
  42. require.NotContains(t, fmt.Sprintf("%+v", billing.captures[0]), batchImageTestData)
  43. require.NotContains(t, fmt.Sprintf("%+v", billing.captures[0]), "gs://")
  44. require.NotContains(t, fmt.Sprintf("%+v", billing.captures[0]), "prompt")
  45. }
  46. func TestBatchImageSettlementService_ZeroSuccessCanComplete(t *testing.T) {
  47. repo := newFakeBatchImageRepository()
  48. job := testSettlingBatchImageJob("imgbatch_zero")
  49. job.SuccessCount = 0
  50. job.FailCount = 4
  51. job.ItemCount = 4
  52. repo.jobs[job.BatchID] = job
  53. billing := &fakeBatchImageBillingRepo{}
  54. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  55. result, err := svc.Settle(context.Background(), job.BatchID)
  56. require.NoError(t, err)
  57. require.Equal(t, 0.0, result.ActualCost)
  58. require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status)
  59. require.Len(t, billing.captures, 1)
  60. require.Equal(t, 0.0, billing.captures[0].ActualAmount)
  61. }
  62. func TestBatchImageSettlementService_CompletedJobReturnsAlreadySettledWithoutBilling(t *testing.T) {
  63. repo := newFakeBatchImageRepository()
  64. job := testSettlingBatchImageJob("imgbatch_done")
  65. job.Status = BatchImageJobStatusCompleted
  66. cost := 0.5
  67. job.ActualCost = &cost
  68. repo.jobs[job.BatchID] = job
  69. billing := &fakeBatchImageBillingRepo{}
  70. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  71. result, err := svc.Settle(context.Background(), job.BatchID)
  72. require.NoError(t, err)
  73. require.True(t, result.AlreadySettled)
  74. require.Equal(t, 0.5, result.ActualCost)
  75. require.Empty(t, billing.captures)
  76. }
  77. func TestBatchImageSettlementService_IdempotentAfterBillingCrash(t *testing.T) {
  78. repo := newFakeBatchImageRepository()
  79. job := testSettlingBatchImageJob("imgbatch_crash")
  80. repo.jobs[job.BatchID] = job
  81. billing := &fakeBatchImageBillingRepo{alreadyApplied: map[string]bool{BatchImageCaptureRequestID(job.BatchID): true}}
  82. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  83. result, err := svc.Settle(context.Background(), job.BatchID)
  84. require.NoError(t, err)
  85. require.Equal(t, 0.5, result.ActualCost)
  86. require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status)
  87. require.Len(t, billing.captures, 1)
  88. }
  89. func TestBatchImageSettlementService_ValidationErrors(t *testing.T) {
  90. tests := []struct {
  91. name string
  92. mutate func(*BatchImageJob)
  93. pricing BatchImagePricingResolver
  94. want error
  95. }{
  96. {name: "invalid_status", mutate: func(j *BatchImageJob) { j.Status = BatchImageJobStatusRunning }, want: ErrBatchImageSettlementInvalidStatus},
  97. {name: "negative_success_count", mutate: func(j *BatchImageJob) { j.SuccessCount = -1 }, want: ErrBatchImageSettlementInvalidCounts},
  98. {name: "negative_fail_count", mutate: func(j *BatchImageJob) { j.FailCount = -1 }, want: ErrBatchImageSettlementInvalidCounts},
  99. {name: "counts_exceed_item_count", mutate: func(j *BatchImageJob) { j.SuccessCount = 2; j.FailCount = 2; j.ItemCount = 3 }, want: ErrBatchImageSettlementInvalidCounts},
  100. {name: "missing_api_key", mutate: func(j *BatchImageJob) { j.APIKeyID = nil }, want: ErrBatchImageSettlementMissingAPIKeyID},
  101. {name: "missing_account", mutate: func(j *BatchImageJob) { j.AccountID = nil }, want: ErrBatchImageSettlementMissingAccountID},
  102. {name: "pricing_missing", pricing: &fakeBatchImagePricingResolver{err: ErrBatchImageSettlementPricingMissing}, want: ErrBatchImageSettlementPricingMissing},
  103. {name: "manifest_conflict", mutate: func(j *BatchImageJob) { v := "different"; j.ManifestHash = &v }, want: ErrBatchImageSettlementManifestConflict},
  104. }
  105. for _, tt := range tests {
  106. t.Run(tt.name, func(t *testing.T) {
  107. repo := newFakeBatchImageRepository()
  108. job := testSettlingBatchImageJob("imgbatch_" + tt.name)
  109. if tt.mutate != nil {
  110. tt.mutate(job)
  111. }
  112. repo.jobs[job.BatchID] = job
  113. pricing := tt.pricing
  114. if pricing == nil {
  115. pricing = &fakeBatchImagePricingResolver{unitPrice: 0.25}
  116. }
  117. billing := &fakeBatchImageBillingRepo{}
  118. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: pricing}
  119. _, err := svc.Settle(context.Background(), job.BatchID)
  120. require.ErrorIs(t, err, tt.want)
  121. require.Empty(t, billing.captures)
  122. require.NotEqual(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status)
  123. })
  124. }
  125. }
  126. func TestBatchImageSettlementService_CostExceedingHoldDoesNotCharge(t *testing.T) {
  127. repo := newFakeBatchImageRepository()
  128. job := testSettlingBatchImageJob("imgbatch_cost_over_hold")
  129. job.SuccessCount = 2
  130. job.FailCount = 0
  131. job.ItemCount = 2
  132. holdAmount := 0.5
  133. job.HoldAmount = &holdAmount
  134. job.EstimatedCost = holdAmount
  135. repo.jobs[job.BatchID] = job
  136. billing := &fakeBatchImageBillingRepo{}
  137. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}}
  138. _, err := svc.Settle(context.Background(), job.BatchID)
  139. require.ErrorIs(t, err, ErrBatchImageSettlementCostExceedsHold)
  140. require.Empty(t, billing.captures)
  141. require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status)
  142. require.Equal(t, "SETTLEMENT_COST_EXCEEDS_HOLD", batchImageDerefString(repo.jobs[job.BatchID].LastErrorCode))
  143. }
  144. func TestBatchImageSettlementService_UsesSubmittedPricingSnapshot(t *testing.T) {
  145. repo := newFakeBatchImageRepository()
  146. job := testSettlingBatchImageJob("imgbatch_snapshot")
  147. job.SuccessCount = 2
  148. job.FailCount = 0
  149. job.ItemCount = 2
  150. job.PricingSnapshotVersion = 1
  151. job.BaseUnitPrice = 0.25
  152. job.GroupRateMultiplier = 1
  153. job.AccountRateMultiplier = 1
  154. job.BatchDiscountMultiplier = 1
  155. job.HoldMultiplier = 1.1
  156. job.BillableUnitPrice = 0.25
  157. job.HoldUnitPrice = 0.275
  158. holdAmount := 0.55
  159. job.HoldAmount = &holdAmount
  160. job.EstimatedCost = 0.5
  161. repo.jobs[job.BatchID] = job
  162. billing := &fakeBatchImageBillingRepo{}
  163. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}}
  164. result, err := svc.Settle(context.Background(), job.BatchID)
  165. require.NoError(t, err)
  166. require.InDelta(t, 0.5, result.ActualCost, 1e-12)
  167. require.Len(t, billing.captures, 1)
  168. require.InDelta(t, 0.5, billing.captures[0].ActualAmount, 1e-12)
  169. require.InDelta(t, 0.55, billing.captures[0].HoldAmount, 1e-12)
  170. }
  171. func TestBatchImageSettlementService_BillingFailureLeavesSettlingAndRecordsError(t *testing.T) {
  172. repo := newFakeBatchImageRepository()
  173. job := testSettlingBatchImageJob("imgbatch_billing_fail")
  174. repo.jobs[job.BatchID] = job
  175. billing := &fakeBatchImageBillingRepo{err: errors.New("temporary billing timeout with gs://hidden-output")}
  176. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  177. _, err := svc.Settle(context.Background(), job.BatchID)
  178. require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed)
  179. require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status)
  180. require.Equal(t, "SETTLEMENT_BILLING_FAILED", batchImageDerefString(repo.jobs[job.BatchID].LastErrorCode))
  181. require.Contains(t, batchImageDerefString(repo.jobs[job.BatchID].LastErrorMessage), "temporary billing timeout")
  182. require.NotNil(t, billing.captures[0])
  183. }
  184. func TestBatchImagePipelineProcessor_SettlesQueuedSettlingJob(t *testing.T) {
  185. repo := newFakeBatchImageRepository()
  186. job := testSettlingBatchImageJob("imgbatch_pipeline")
  187. repo.jobs[job.BatchID] = job
  188. billing := &fakeBatchImageBillingRepo{}
  189. settlement := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  190. processor := &BatchImagePipelineProcessor{
  191. ProviderProcessor: &BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}},
  192. SettlementService: settlement,
  193. }
  194. result, err := processor.Process(context.Background(), job.BatchID)
  195. require.NoError(t, err)
  196. require.True(t, result.Terminal)
  197. require.Equal(t, BatchImageJobStatusCompleted, repo.jobs[job.BatchID].Status)
  198. require.Len(t, billing.captures, 1)
  199. }
  200. func TestBatchImagePipelineProcessor_RequeuesTransientSettlementFailure(t *testing.T) {
  201. repo := newFakeBatchImageRepository()
  202. job := testSettlingBatchImageJob("imgbatch_pipeline_retry")
  203. repo.jobs[job.BatchID] = job
  204. settlement := &BatchImageSettlementService{Repo: repo, BillingRepo: &fakeBatchImageBillingRepo{err: errors.New("temporary")}, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  205. processor := &BatchImagePipelineProcessor{
  206. ProviderProcessor: &BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}},
  207. SettlementService: settlement,
  208. }
  209. result, err := processor.Process(context.Background(), job.BatchID)
  210. require.NoError(t, err)
  211. require.False(t, result.Terminal)
  212. require.Equal(t, batchImageSettlementRetryDelay, result.RequeueAfter)
  213. require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status)
  214. }
  215. func TestBatchImagePipelineProcessor_FailsAndReleasesAfterSettlementRetryLimit(t *testing.T) {
  216. repo := newFakeBatchImageRepository()
  217. job := testSettlingBatchImageJob("imgbatch_pipeline_retry_exhausted")
  218. job.RetryCount = batchImageSettlementMaxRetries - 1
  219. repo.jobs[job.BatchID] = job
  220. billing := &fakeBatchImageBillingRepo{captureErr: errors.New("temporary billing timeout")}
  221. settlement := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  222. processor := &BatchImagePipelineProcessor{
  223. ProviderProcessor: &BatchImageProviderProcessor{Repo: repo, ProviderRegistry: NewBatchImageProviderRegistry(&fakeProcessorProvider{}), AccountResolver: &fakeBatchImageAccountResolver{account: &Account{}}},
  224. SettlementService: settlement,
  225. }
  226. result, err := processor.Process(context.Background(), job.BatchID)
  227. require.NoError(t, err)
  228. require.True(t, result.Terminal)
  229. require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status)
  230. require.Equal(t, "SETTLEMENT_BILLING_RETRY_EXHAUSTED", batchImageDerefString(repo.jobs[job.BatchID].LastErrorCode))
  231. require.Len(t, billing.captures, 1)
  232. require.Len(t, billing.releases, 1)
  233. require.Equal(t, BatchImageReleaseRequestID(job.BatchID), billing.releases[0].RequestID)
  234. }
  235. func TestBatchImageSettlementRetryExhaustedReleaseIsIdempotentAfterTransitionFailure(t *testing.T) {
  236. repo := newFakeBatchImageRepository()
  237. job := testSettlingBatchImageJob("imgbatch_retry_exhausted_transition_fail")
  238. job.RetryCount = batchImageSettlementMaxRetries
  239. job.LastErrorCode = batchImageStringPtr("SETTLEMENT_BILLING_FAILED")
  240. repo.jobs[job.BatchID] = job
  241. repo.transitionErr = errors.New("temporary transition failure")
  242. billing := &fakeBatchImageBillingRepo{}
  243. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  244. _, err := svc.Settle(context.Background(), job.BatchID)
  245. require.ErrorContains(t, err, "temporary transition failure")
  246. require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status)
  247. require.Len(t, billing.releases, 1)
  248. require.Len(t, billing.seen, 1)
  249. repo.transitionErr = nil
  250. _, err = svc.Settle(context.Background(), job.BatchID)
  251. require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed)
  252. require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status)
  253. require.Len(t, billing.releases, 2)
  254. require.Equal(t, billing.releases[0].RequestID, billing.releases[1].RequestID)
  255. require.Len(t, billing.seen, 1)
  256. }
  257. func TestBatchImageSettlementService_CostExceedsHoldExhaustsAndReleases(t *testing.T) {
  258. repo := newFakeBatchImageRepository()
  259. job := testSettlingBatchImageJob("imgbatch_over_hold_exhausted")
  260. job.SuccessCount = 2
  261. job.FailCount = 0
  262. job.ItemCount = 2
  263. holdAmount := 0.5
  264. job.HoldAmount = &holdAmount
  265. job.EstimatedCost = holdAmount
  266. requestHash := "request-hash-over-hold"
  267. job.RequestHash = &requestHash
  268. repo.jobs[job.BatchID] = job
  269. billing := &fakeBatchImageBillingRepo{}
  270. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}}
  271. // 前 N-1 次:记录失败并返回错误(等待 worker 重试)。
  272. for i := 0; i < batchImageSettlementMaxRetries-1; i++ {
  273. _, err := svc.Settle(context.Background(), job.BatchID)
  274. require.ErrorIs(t, err, ErrBatchImageSettlementCostExceedsHold)
  275. require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status)
  276. }
  277. // 达到上限:必须走耗尽出口释放冻结并转 failed,而不是无限 requeue。
  278. _, err := svc.Settle(context.Background(), job.BatchID)
  279. require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed)
  280. require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status)
  281. require.Empty(t, billing.captures)
  282. require.Len(t, billing.releases, 1)
  283. require.Equal(t, BatchImageReleaseRequestID(job.BatchID), billing.releases[0].RequestID)
  284. // 释放指纹必须与 processor/Cancel/recovery 一致地使用 RequestHash,
  285. // 否则共享同一 request id 的后续释放会命中指纹冲突(毒消息)。
  286. require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash)
  287. }
  288. func TestBatchImageSettlementService_InvalidCountsExhaustsAndReleases(t *testing.T) {
  289. repo := newFakeBatchImageRepository()
  290. job := testSettlingBatchImageJob("imgbatch_bad_counts_exhausted")
  291. job.SuccessCount = 2
  292. job.FailCount = 2
  293. job.ItemCount = 3
  294. requestHash := "request-hash-bad-counts"
  295. job.RequestHash = &requestHash
  296. repo.jobs[job.BatchID] = job
  297. billing := &fakeBatchImageBillingRepo{}
  298. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  299. for i := 0; i < batchImageSettlementMaxRetries-1; i++ {
  300. _, err := svc.Settle(context.Background(), job.BatchID)
  301. require.ErrorIs(t, err, ErrBatchImageSettlementInvalidCounts)
  302. }
  303. _, err := svc.Settle(context.Background(), job.BatchID)
  304. require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed)
  305. require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status)
  306. require.Empty(t, billing.captures)
  307. require.Len(t, billing.releases, 1)
  308. require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash)
  309. }
  310. func TestReleaseBatchImageBalanceHold_TreatsFingerprintConflictAsReleased(t *testing.T) {
  311. job := testSettlingBatchImageJob("imgbatch_release_conflict")
  312. // 历史版本用 manifestHash 释放过一次:同一 request id 再以 RequestHash
  313. // 释放会命中指纹冲突。资金已归还,必须视为幂等成功而非毒消息。
  314. billing := &fakeBatchImageBillingRepo{releaseErr: ErrUsageBillingRequestConflict}
  315. err := releaseBatchImageBalanceHold(context.Background(), billing, job, "request-hash")
  316. require.NoError(t, err)
  317. require.Len(t, billing.releases, 1)
  318. }
  319. func TestBatchImageSettlementManifestHash(t *testing.T) {
  320. job := testSettlingBatchImageJob("imgbatch_hash")
  321. first := BuildBatchImageSettlementManifestHash(job)
  322. job.CreatedAt = job.CreatedAt.AddDate(0, 0, 1)
  323. job.UpdatedAt = job.UpdatedAt.AddDate(0, 0, 1)
  324. require.Equal(t, first, BuildBatchImageSettlementManifestHash(job))
  325. job.SuccessCount++
  326. require.NotEqual(t, first, BuildBatchImageSettlementManifestHash(job))
  327. job.SuccessCount--
  328. promptOrBase64 := first + " prompt " + batchImageTestData
  329. require.NotContains(t, BuildBatchImageSettlementManifestHash(job), promptOrBase64)
  330. }
  331. func TestBatchImageSettlementBillingRequestIDs(t *testing.T) {
  332. repo := newFakeBatchImageRepository()
  333. first := testSettlingBatchImageJob("imgbatch_unique_1")
  334. second := testSettlingBatchImageJob("imgbatch_unique_2")
  335. repo.jobs[first.BatchID] = first
  336. repo.jobs[second.BatchID] = second
  337. billing := &fakeBatchImageBillingRepo{}
  338. svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
  339. _, err := svc.Settle(context.Background(), first.BatchID)
  340. require.NoError(t, err)
  341. _, err = svc.Settle(context.Background(), first.BatchID)
  342. require.NoError(t, err)
  343. _, err = svc.Settle(context.Background(), second.BatchID)
  344. require.NoError(t, err)
  345. require.Len(t, billing.captures, 2)
  346. require.Equal(t, BatchImageCaptureRequestID(first.BatchID), billing.captures[0].RequestID)
  347. require.Equal(t, BatchImageCaptureRequestID(second.BatchID), billing.captures[1].RequestID)
  348. require.NotEqual(t, billing.captures[0].RequestID, billing.captures[1].RequestID)
  349. require.Len(t, billing.seen, 2)
  350. }
  351. func testSettlingBatchImageJob(batchID string) *BatchImageJob {
  352. apiKeyID := int64(321)
  353. accountID := int64(654)
  354. providerJobName := "providers/job"
  355. outputRef := "files/output"
  356. holdAmount := 1.25
  357. holdID := BatchImageHoldRequestID(batchID)
  358. return &BatchImageJob{
  359. BatchID: batchID,
  360. UserID: 123,
  361. APIKeyID: &apiKeyID,
  362. AccountID: &accountID,
  363. Provider: BatchImageProviderGeminiAPI,
  364. Model: "gemini-image",
  365. Status: BatchImageJobStatusSettling,
  366. ProviderJobName: &providerJobName,
  367. ProviderOutputRef: &outputRef,
  368. ItemCount: 3,
  369. SuccessCount: 2,
  370. FailCount: 1,
  371. EstimatedCost: holdAmount,
  372. HoldAmount: &holdAmount,
  373. HoldID: &holdID,
  374. }
  375. }
  376. type fakeBatchImagePricingResolver struct {
  377. unitPrice float64
  378. missingModels map[string]bool
  379. err error
  380. }
  381. func (r *fakeBatchImagePricingResolver) BatchImageUnitPrice(_ context.Context, job *BatchImageJob) (float64, error) {
  382. if r.err != nil {
  383. return 0, r.err
  384. }
  385. if job != nil && r.missingModels[job.Model] {
  386. return 0, ErrBatchImageSettlementPricingMissing
  387. }
  388. return r.unitPrice, nil
  389. }
  390. type fakeBatchImageBillingRepo struct {
  391. commands []*UsageBillingCommand
  392. reserves []*BatchImageBalanceHoldCommand
  393. captures []*BatchImageBalanceHoldCommand
  394. releases []*BatchImageBalanceHoldCommand
  395. seen map[string]struct{}
  396. alreadyApplied map[string]bool
  397. err error
  398. reserveErr error
  399. captureErr error
  400. releaseErr error
  401. }
  402. func (r *fakeBatchImageBillingRepo) Apply(_ context.Context, cmd *UsageBillingCommand) (*UsageBillingApplyResult, error) {
  403. if r.seen == nil {
  404. r.seen = make(map[string]struct{})
  405. }
  406. if r.err != nil {
  407. r.commands = append(r.commands, cmd)
  408. return nil, r.err
  409. }
  410. if cmd != nil {
  411. cmd.Normalize()
  412. if _, ok := r.seen[cmd.RequestID]; ok || r.alreadyApplied[cmd.RequestID] {
  413. r.commands = append(r.commands, cmd)
  414. return &UsageBillingApplyResult{Applied: false}, nil
  415. }
  416. r.seen[cmd.RequestID] = struct{}{}
  417. }
  418. r.commands = append(r.commands, cmd)
  419. return &UsageBillingApplyResult{Applied: true}, nil
  420. }
  421. func (r *fakeBatchImageBillingRepo) ReserveBatchImageBalance(_ context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) {
  422. if r.reserveErr != nil {
  423. r.reserves = append(r.reserves, cmd)
  424. return nil, r.reserveErr
  425. }
  426. return r.applyHold(cmd, &r.reserves)
  427. }
  428. func (r *fakeBatchImageBillingRepo) CaptureBatchImageBalance(_ context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) {
  429. if r.captureErr != nil {
  430. r.captures = append(r.captures, cmd)
  431. return nil, r.captureErr
  432. }
  433. return r.applyHold(cmd, &r.captures)
  434. }
  435. func (r *fakeBatchImageBillingRepo) ReleaseBatchImageBalance(_ context.Context, cmd *BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) {
  436. if r.releaseErr != nil {
  437. r.releases = append(r.releases, cmd)
  438. return nil, r.releaseErr
  439. }
  440. return r.applyHold(cmd, &r.releases)
  441. }
  442. func (r *fakeBatchImageBillingRepo) applyHold(cmd *BatchImageBalanceHoldCommand, calls *[]*BatchImageBalanceHoldCommand) (*BatchImageBalanceHoldResult, error) {
  443. if r.seen == nil {
  444. r.seen = make(map[string]struct{})
  445. }
  446. if r.err != nil {
  447. *calls = append(*calls, cmd)
  448. return nil, r.err
  449. }
  450. if cmd != nil {
  451. cmd.Normalize()
  452. if _, ok := r.seen[cmd.RequestID]; ok || r.alreadyApplied[cmd.RequestID] {
  453. *calls = append(*calls, cmd)
  454. return &BatchImageBalanceHoldResult{Applied: false}, nil
  455. }
  456. r.seen[cmd.RequestID] = struct{}{}
  457. }
  458. *calls = append(*calls, cmd)
  459. return &BatchImageBalanceHoldResult{Applied: true}, nil
  460. }
  461. var _ UsageBillingRepository = (*fakeBatchImageBillingRepo)(nil)
  462. var _ BatchImagePricingResolver = (*fakeBatchImagePricingResolver)(nil)
  463. var _ = strings.TrimSpace