You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

139 line
4.6 KiB

  1. package service
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/common"
  5. "github.com/QuantumNous/new-api/model"
  6. "github.com/QuantumNous/new-api/setting/system_setting"
  7. "github.com/glebarez/sqlite"
  8. "github.com/stretchr/testify/assert"
  9. "github.com/stretchr/testify/require"
  10. "gorm.io/gorm"
  11. )
  12. func setupBillingSyncDB(t *testing.T) *gorm.DB {
  13. t.Helper()
  14. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  15. require.NoError(t, err)
  16. sqlDB, _ := db.DB()
  17. sqlDB.SetMaxOpenConns(1)
  18. origDB := model.DB
  19. origLogDB := model.LOG_DB
  20. model.DB = db
  21. model.LOG_DB = db
  22. common.UsingSQLite = true
  23. common.RedisEnabled = false
  24. require.NoError(t, db.AutoMigrate(&model.PendingSyncRecord{}, &model.User{}))
  25. t.Cleanup(func() {
  26. model.DB = origDB
  27. model.LOG_DB = origLogDB
  28. sqlDB.Close()
  29. })
  30. return db
  31. }
  32. // TestNeedsRefundLocked_SyncedUser 测试同步用户在 consumed > 0 时需要退款
  33. func TestNeedsRefundLocked_SyncedUser(t *testing.T) {
  34. setupBillingSyncDB(t)
  35. settings := system_setting.GetRegionSyncSettings()
  36. origThreshold := settings.MinBalanceThreshold
  37. defer func() { settings.MinBalanceThreshold = origThreshold }()
  38. settings.MinBalanceThreshold = 100000
  39. require.NoError(t, model.DB.Create(&model.User{Id: 100, Username: "test-nr", Source: common.UserSourceSynced, SyncedQuota: 500000}).Error)
  40. f := NewSyncedUserFunding(100, 200, "req-nr-001", 500000)
  41. require.NoError(t, f.PreConsume(5000))
  42. session := &BillingSession{
  43. funding: f,
  44. }
  45. assert.True(t, session.needsRefundLocked())
  46. }
  47. // TestNeedsRefundLocked_SyncedUser_NoConsumed 测试同步用户在 consumed == 0 时不需要退款
  48. func TestNeedsRefundLocked_SyncedUser_NoConsumed(t *testing.T) {
  49. setupBillingSyncDB(t)
  50. f := NewSyncedUserFunding(100, 200, "req-nr-002", 500000)
  51. session := &BillingSession{
  52. funding: f,
  53. }
  54. assert.False(t, session.needsRefundLocked())
  55. }
  56. // TestNeedsRefundLocked_SyncedUser_FundingSettled 测试 fundingSettled=true 后不需要退款
  57. func TestNeedsRefundLocked_SyncedUser_FundingSettled(t *testing.T) {
  58. setupBillingSyncDB(t)
  59. settings := system_setting.GetRegionSyncSettings()
  60. origThreshold := settings.MinBalanceThreshold
  61. defer func() { settings.MinBalanceThreshold = origThreshold }()
  62. settings.MinBalanceThreshold = 100000
  63. require.NoError(t, model.DB.Create(&model.User{Id: 100, Username: "test-nr3", Source: common.UserSourceSynced, SyncedQuota: 500000}).Error)
  64. f := NewSyncedUserFunding(100, 200, "req-nr-003", 500000)
  65. require.NoError(t, f.PreConsume(5000))
  66. session := &BillingSession{
  67. funding: f,
  68. fundingSettled: true, // 资金已提交
  69. }
  70. // fundingSettled=true 时不应再退
  71. assert.False(t, session.needsRefundLocked())
  72. }
  73. // TestShouldTrust_SyncedUser_SourceNotWalletOrSub 验证同步用户 Source 不是 wallet/subscription,
  74. // 因此在 shouldTrust 的 switch 中走 default 分支返回 false
  75. func TestShouldTrust_SyncedUser_SourceNotWalletOrSub(t *testing.T) {
  76. f := NewSyncedUserFunding(100, 200, "req-trust", 500000)
  77. assert.Equal(t, "synced_wallet", f.Source())
  78. assert.NotEqual(t, BillingSourceWallet, f.Source())
  79. assert.NotEqual(t, BillingSourceSubscription, f.Source())
  80. }
  81. // TestSyncedUserFunding_BillingSource 测试同步用户计费来源标识
  82. func TestSyncedUserFunding_BillingSource(t *testing.T) {
  83. f := NewSyncedUserFunding(100, 200, "req-src", 500000)
  84. assert.Equal(t, BillingSourceSyncedWallet, f.Source())
  85. assert.Equal(t, "synced_wallet", f.Source())
  86. }
  87. // TestNeedsRefundLocked_WalletAndSynced 对比钱包和同步用户的行为
  88. func TestNeedsRefundLocked_WalletAndSynced(t *testing.T) {
  89. // 钱包用户:tokenConsumed > 0 才需要退
  90. walletSession := &BillingSession{
  91. funding: &WalletFunding{userId: 1, consumed: 100},
  92. tokenConsumed: 100,
  93. }
  94. assert.True(t, walletSession.needsRefundLocked())
  95. walletSessionZero := &BillingSession{
  96. funding: &WalletFunding{userId: 1, consumed: 100},
  97. tokenConsumed: 0,
  98. }
  99. assert.False(t, walletSessionZero.needsRefundLocked())
  100. // 同步用户:consumed > 0 就需要退(即使 tokenConsumed=0)
  101. setupBillingSyncDB(t)
  102. settings := system_setting.GetRegionSyncSettings()
  103. origThreshold := settings.MinBalanceThreshold
  104. defer func() { settings.MinBalanceThreshold = origThreshold }()
  105. settings.MinBalanceThreshold = 100000
  106. require.NoError(t, model.DB.Create(&model.User{Id: 100, Username: "test-compare", Source: common.UserSourceSynced, SyncedQuota: 500000}).Error)
  107. syncF := NewSyncedUserFunding(100, 200, "req-compare", 500000)
  108. require.NoError(t, syncF.PreConsume(5000))
  109. syncSession := &BillingSession{
  110. funding: syncF,
  111. tokenConsumed: 0, // tokenConsumed=0 但 consumed > 0
  112. }
  113. assert.True(t, syncSession.needsRefundLocked())
  114. }