batch_image_billing_recovery_test.go 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "time"
  8. "github.com/stretchr/testify/require"
  9. )
  10. type recordingBatchImageQueue struct {
  11. *fakeBatchImageQueue
  12. enqueued []string
  13. }
  14. func (q *recordingBatchImageQueue) Enqueue(_ context.Context, batchID string) error {
  15. q.enqueued = append(q.enqueued, batchID)
  16. return nil
  17. }
  18. func TestBatchImageBillingRecoveryService_ReleasesStaleUnsubmittedHold(t *testing.T) {
  19. repo := newFakeBatchImageRepository()
  20. apiKeyID := int64(22)
  21. holdAmount := 0.5
  22. stale := &BatchImageJob{
  23. BatchID: "imgbatch_stale_created",
  24. UserID: 11,
  25. APIKeyID: &apiKeyID,
  26. Status: BatchImageJobStatusCreated,
  27. EstimatedCost: holdAmount,
  28. HoldAmount: &holdAmount,
  29. CreatedAt: time.Now().Add(-time.Hour),
  30. UpdatedAt: time.Now().Add(-time.Hour),
  31. }
  32. activeProviderName := "providers/job"
  33. active := &BatchImageJob{
  34. BatchID: "imgbatch_has_provider",
  35. UserID: 11,
  36. APIKeyID: &apiKeyID,
  37. Status: BatchImageJobStatusSubmitted,
  38. ProviderJobName: &activeProviderName,
  39. EstimatedCost: holdAmount,
  40. HoldAmount: &holdAmount,
  41. CreatedAt: time.Now().Add(-time.Hour),
  42. UpdatedAt: time.Now().Add(-time.Hour),
  43. }
  44. repo.jobs[stale.BatchID] = stale
  45. repo.jobs[active.BatchID] = active
  46. billing := &fakeBatchImageBillingRepo{}
  47. svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, StaleAfter: time.Minute, Limit: 10}
  48. released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background())
  49. require.NoError(t, err)
  50. require.Equal(t, 1, released)
  51. require.Equal(t, BatchImageJobStatusFailed, repo.jobs[stale.BatchID].Status)
  52. require.Equal(t, "SUBMIT_STALE_BEFORE_PROVIDER", batchImageDerefString(repo.jobs[stale.BatchID].LastErrorCode))
  53. require.Len(t, billing.releases, 1)
  54. require.Equal(t, BatchImageReleaseRequestID(stale.BatchID), billing.releases[0].RequestID)
  55. require.Equal(t, BatchImageJobStatusSubmitted, repo.jobs[active.BatchID].Status)
  56. }
  57. func TestBatchImageBillingRecoveryService_SkipsJobRefreshedByHeartbeat(t *testing.T) {
  58. repo := newFakeBatchImageRepository()
  59. apiKeyID := int64(22)
  60. holdAmount := 0.5
  61. // updated_at 在 cutoff 之后(慢提交心跳持续续期):不得误杀退款。
  62. fresh := &BatchImageJob{
  63. BatchID: "imgbatch_fresh_uploading",
  64. UserID: 11,
  65. APIKeyID: &apiKeyID,
  66. Status: BatchImageJobStatusUploading,
  67. EstimatedCost: holdAmount,
  68. HoldAmount: &holdAmount,
  69. CreatedAt: time.Now().Add(-time.Hour),
  70. UpdatedAt: time.Now(),
  71. }
  72. repo.jobs[fresh.BatchID] = fresh
  73. billing := &fakeBatchImageBillingRepo{}
  74. svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, StaleAfter: time.Minute, Limit: 10}
  75. released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background())
  76. require.NoError(t, err)
  77. require.Equal(t, 0, released)
  78. require.Equal(t, BatchImageJobStatusUploading, repo.jobs[fresh.BatchID].Status)
  79. require.Empty(t, billing.releases)
  80. }
  81. func TestBatchImageBillingRecoveryService_EnqueuesRetryWhenReleaseFails(t *testing.T) {
  82. repo := newFakeBatchImageRepository()
  83. apiKeyID := int64(22)
  84. holdAmount := 0.5
  85. stale := &BatchImageJob{
  86. BatchID: "imgbatch_stale_release_fail",
  87. UserID: 11,
  88. APIKeyID: &apiKeyID,
  89. Status: BatchImageJobStatusCreated,
  90. EstimatedCost: holdAmount,
  91. HoldAmount: &holdAmount,
  92. CreatedAt: time.Now().Add(-time.Hour),
  93. UpdatedAt: time.Now().Add(-time.Hour),
  94. }
  95. repo.jobs[stale.BatchID] = stale
  96. billing := &fakeBatchImageBillingRepo{releaseErr: errors.New("billing db down")}
  97. queue := &recordingBatchImageQueue{fakeBatchImageQueue: newFakeBatchImageQueue("")}
  98. svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, Queue: queue, StaleAfter: time.Minute, Limit: 10}
  99. released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background())
  100. // job 已转 failed、不会再出现在 stale 列表:释放失败必须入队重试
  101. //(由 worker 的 releaseTerminalHold 兜底),否则冻结余额永久泄漏。
  102. require.Error(t, err)
  103. require.Equal(t, 0, released)
  104. require.Equal(t, BatchImageJobStatusFailed, repo.jobs[stale.BatchID].Status)
  105. require.Equal(t, []string{stale.BatchID}, queue.enqueued)
  106. }