Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

255 linhas
7.4 KiB

  1. package service
  2. import (
  3. "testing"
  4. "time"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/model"
  7. "github.com/QuantumNous/new-api/setting/system_setting"
  8. "github.com/glebarez/sqlite"
  9. "github.com/stretchr/testify/assert"
  10. "github.com/stretchr/testify/require"
  11. "gorm.io/gorm"
  12. )
  13. func setupSyncedFundingDB(t *testing.T) *gorm.DB {
  14. t.Helper()
  15. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  16. require.NoError(t, err)
  17. sqlDB, _ := db.DB()
  18. sqlDB.SetMaxOpenConns(1)
  19. origDB := model.DB
  20. origLogDB := model.LOG_DB
  21. model.DB = db
  22. model.LOG_DB = db
  23. common.UsingSQLite = true
  24. common.RedisEnabled = false
  25. require.NoError(t, db.AutoMigrate(&model.PendingSyncRecord{}, &model.User{}))
  26. t.Cleanup(func() {
  27. model.DB = origDB
  28. model.LOG_DB = origLogDB
  29. sqlDB.Close()
  30. })
  31. return db
  32. }
  33. func TestNewSyncedUserFunding(t *testing.T) {
  34. f := NewSyncedUserFunding(100, 200, "req-001", 500000)
  35. assert.Equal(t, 100, f.userId)
  36. assert.Equal(t, 200, f.remoteUserId)
  37. assert.Equal(t, "req-001", f.requestId)
  38. assert.Equal(t, 500000, f.syncedQuota)
  39. assert.Equal(t, 0, f.consumed)
  40. }
  41. func TestSyncedUserFunding_Source(t *testing.T) {
  42. f := NewSyncedUserFunding(100, 200, "req-001", 500000)
  43. assert.Equal(t, BillingSourceSyncedWallet, f.Source())
  44. }
  45. func TestPreConsume_Success(t *testing.T) {
  46. setupSyncedFundingDB(t)
  47. settings := system_setting.GetRegionSyncSettings()
  48. origThreshold := settings.MinBalanceThreshold
  49. settings.MinBalanceThreshold = 100000
  50. defer func() { settings.MinBalanceThreshold = origThreshold }()
  51. // 创建测试用户(AtomicDecreaseSyncedQuota 操作 DB 中的 synced_quota)
  52. require.NoError(t, model.DB.Create(&model.User{
  53. Id: 100, Username: "test-synced", Source: common.UserSourceSynced, SyncedQuota: 1000000,
  54. }).Error)
  55. f := NewSyncedUserFunding(100, 200, "req-pre-success", 1000000)
  56. err := f.PreConsume(50000)
  57. require.NoError(t, err)
  58. assert.Equal(t, 50000, f.consumed)
  59. // 验证 PendingSyncRecord 已创建
  60. var record model.PendingSyncRecord
  61. require.NoError(t, model.DB.Where("request_id = ?", "req-pre-success").First(&record).Error)
  62. assert.Equal(t, 100, record.UserId)
  63. assert.Equal(t, 200, record.RemoteUserId)
  64. assert.Equal(t, 50000, record.Quota)
  65. assert.Equal(t, model.PendingSyncStatusPending, record.Status)
  66. }
  67. func TestPreConsume_ZeroAmount(t *testing.T) {
  68. setupSyncedFundingDB(t)
  69. f := NewSyncedUserFunding(100, 200, "req-pre-zero", 1000000)
  70. err := f.PreConsume(0)
  71. assert.NoError(t, err)
  72. assert.Equal(t, 0, f.consumed)
  73. }
  74. func TestPreConsume_NegativeAmount(t *testing.T) {
  75. setupSyncedFundingDB(t)
  76. f := NewSyncedUserFunding(100, 200, "req-pre-neg", 1000000)
  77. err := f.PreConsume(-100)
  78. assert.NoError(t, err)
  79. assert.Equal(t, 0, f.consumed)
  80. }
  81. func TestPreConsume_InsufficientQuota(t *testing.T) {
  82. setupSyncedFundingDB(t)
  83. settings := system_setting.GetRegionSyncSettings()
  84. origThreshold := settings.MinBalanceThreshold
  85. settings.MinBalanceThreshold = 100000
  86. defer func() { settings.MinBalanceThreshold = origThreshold }()
  87. require.NoError(t, model.DB.Create(&model.User{
  88. Id: 100, Username: "test-insuf", Source: common.UserSourceSynced, SyncedQuota: 149999,
  89. }).Error)
  90. // syncedQuota=149999, amount=50000, threshold=100000
  91. // 149999 < 50000 + 100000 = 150000 → 不满足
  92. f := NewSyncedUserFunding(100, 200, "req-pre-insuf", 149999)
  93. err := f.PreConsume(50000)
  94. assert.Error(t, err)
  95. assert.Contains(t, err.Error(), "余额不足")
  96. assert.Equal(t, 0, f.consumed)
  97. }
  98. func TestPreConsume_ExactlyAtThreshold(t *testing.T) {
  99. setupSyncedFundingDB(t)
  100. settings := system_setting.GetRegionSyncSettings()
  101. origThreshold := settings.MinBalanceThreshold
  102. settings.MinBalanceThreshold = 100000
  103. defer func() { settings.MinBalanceThreshold = origThreshold }()
  104. require.NoError(t, model.DB.Create(&model.User{
  105. Id: 100, Username: "test-exact", Source: common.UserSourceSynced, SyncedQuota: 150000,
  106. }).Error)
  107. // syncedQuota=150000, amount=50000, threshold=100000
  108. // 150000 >= 50000 + 100000 = 150000 → 刚好满足
  109. f := NewSyncedUserFunding(100, 200, "req-pre-exact", 150000)
  110. err := f.PreConsume(50000)
  111. assert.NoError(t, err)
  112. assert.Equal(t, 50000, f.consumed)
  113. }
  114. func TestSettle_PositiveDelta(t *testing.T) {
  115. setupSyncedFundingDB(t)
  116. f := NewSyncedUserFunding(100, 200, "req-settle-pos", 500000)
  117. f.consumed = 1000
  118. err := f.Settle(500)
  119. assert.NoError(t, err)
  120. assert.Equal(t, 1500, f.consumed)
  121. assert.Equal(t, 499500, f.syncedQuota) // 500000 - 500
  122. }
  123. func TestSettle_NegativeDelta(t *testing.T) {
  124. setupSyncedFundingDB(t)
  125. f := NewSyncedUserFunding(100, 200, "req-settle-neg", 500000)
  126. f.consumed = 1000
  127. err := f.Settle(-300)
  128. assert.NoError(t, err)
  129. assert.Equal(t, 700, f.consumed)
  130. assert.Equal(t, 500300, f.syncedQuota) // 500000 - (-300) = 500000 + 300
  131. }
  132. func TestSettle_ZeroDelta(t *testing.T) {
  133. f := NewSyncedUserFunding(100, 200, "req-settle-zero", 500000)
  134. f.consumed = 1000
  135. err := f.Settle(0)
  136. assert.NoError(t, err)
  137. assert.Equal(t, 1000, f.consumed) // 不变
  138. assert.Equal(t, 500000, f.syncedQuota) // 不变
  139. }
  140. func TestSettle_TransitionsPendingToSettled(t *testing.T) {
  141. db := setupSyncedFundingDB(t)
  142. // 先创建 pending 记录
  143. require.NoError(t, model.CreatePendingSyncRecord(100, 200, "req-settle-trans", 1000, 1000))
  144. f := NewSyncedUserFunding(100, 200, "req-settle-trans", 500000)
  145. f.consumed = 1000
  146. err := f.Settle(-300) // 实际消费 700,退还 300
  147. assert.NoError(t, err)
  148. assert.Equal(t, 700, f.consumed)
  149. // 验证记录状态从 pending 转为 settled,quota 更新为实际消费
  150. var record model.PendingSyncRecord
  151. require.NoError(t, db.Where("request_id = ?", "req-settle-trans").First(&record).Error)
  152. assert.Equal(t, model.PendingSyncStatusSettled, record.Status)
  153. assert.Equal(t, 700, record.Quota)
  154. }
  155. func TestRefund_Success(t *testing.T) {
  156. db := setupSyncedFundingDB(t)
  157. // 创建测试用户
  158. require.NoError(t, model.DB.Create(&model.User{
  159. Id: 100, Username: "test-refund", Source: common.UserSourceSynced, SyncedQuota: 500000,
  160. }).Error)
  161. // 先创建一条 pending 记录
  162. require.NoError(t, model.CreatePendingSyncRecord(100, 200, "req-refund-ok", 5000, 5000))
  163. f := NewSyncedUserFunding(100, 200, "req-refund-ok", 500000)
  164. f.consumed = 5000
  165. err := f.Refund()
  166. assert.NoError(t, err)
  167. // 验证记录已被标记为 synced
  168. var record model.PendingSyncRecord
  169. require.NoError(t, db.Where("request_id = ?", "req-refund-ok").First(&record).Error)
  170. assert.Equal(t, model.PendingSyncStatusSynced, record.Status)
  171. assert.NotZero(t, record.SyncedAt)
  172. }
  173. func TestRefund_NoConsumed(t *testing.T) {
  174. setupSyncedFundingDB(t)
  175. f := NewSyncedUserFunding(100, 200, "req-refund-nocons", 500000)
  176. f.consumed = 0
  177. err := f.Refund()
  178. assert.NoError(t, err)
  179. }
  180. func TestRefund_RecordNotFound(t *testing.T) {
  181. setupSyncedFundingDB(t)
  182. // 没有 PendingSyncRecord,但 consumed > 0
  183. f := NewSyncedUserFunding(100, 200, "req-refund-notfound", 500000)
  184. f.consumed = 5000
  185. err := f.Refund()
  186. assert.NoError(t, err) // 应优雅处理,不报错
  187. }
  188. func TestRefund_RecordAlreadySynced(t *testing.T) {
  189. db := setupSyncedFundingDB(t)
  190. // 创建已 synced 的记录
  191. require.NoError(t, db.Create(&model.PendingSyncRecord{
  192. UserId: 100, RemoteUserId: 200, RequestId: "req-refund-synced",
  193. Quota: 5000, Status: model.PendingSyncStatusSynced,
  194. SyncedAt: time.Now().Unix(), CreatedAt: time.Now().Unix(),
  195. }).Error)
  196. f := NewSyncedUserFunding(100, 200, "req-refund-synced", 500000)
  197. f.consumed = 5000
  198. err := f.Refund()
  199. assert.NoError(t, err) // status=pending 不匹配,优雅处理
  200. }