Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

136 řádky
4.2 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{}))
  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. f := NewSyncedUserFunding(100, 200, "req-nr-001", 500000)
  40. require.NoError(t, f.PreConsume(5000))
  41. session := &BillingSession{
  42. funding: f,
  43. }
  44. assert.True(t, session.needsRefundLocked())
  45. }
  46. // TestNeedsRefundLocked_SyncedUser_NoConsumed 测试同步用户在 consumed == 0 时不需要退款
  47. func TestNeedsRefundLocked_SyncedUser_NoConsumed(t *testing.T) {
  48. setupBillingSyncDB(t)
  49. f := NewSyncedUserFunding(100, 200, "req-nr-002", 500000)
  50. session := &BillingSession{
  51. funding: f,
  52. }
  53. assert.False(t, session.needsRefundLocked())
  54. }
  55. // TestNeedsRefundLocked_SyncedUser_FundingSettled 测试 fundingSettled=true 后不需要退款
  56. func TestNeedsRefundLocked_SyncedUser_FundingSettled(t *testing.T) {
  57. setupBillingSyncDB(t)
  58. settings := system_setting.GetRegionSyncSettings()
  59. origThreshold := settings.MinBalanceThreshold
  60. defer func() { settings.MinBalanceThreshold = origThreshold }()
  61. settings.MinBalanceThreshold = 100000
  62. f := NewSyncedUserFunding(100, 200, "req-nr-003", 500000)
  63. require.NoError(t, f.PreConsume(5000))
  64. session := &BillingSession{
  65. funding: f,
  66. fundingSettled: true, // 资金已提交
  67. }
  68. // fundingSettled=true 时不应再退
  69. assert.False(t, session.needsRefundLocked())
  70. }
  71. // TestShouldTrust_SyncedUser_SourceNotWalletOrSub 验证同步用户 Source 不是 wallet/subscription,
  72. // 因此在 shouldTrust 的 switch 中走 default 分支返回 false
  73. func TestShouldTrust_SyncedUser_SourceNotWalletOrSub(t *testing.T) {
  74. f := NewSyncedUserFunding(100, 200, "req-trust", 500000)
  75. assert.Equal(t, "synced_wallet", f.Source())
  76. assert.NotEqual(t, BillingSourceWallet, f.Source())
  77. assert.NotEqual(t, BillingSourceSubscription, f.Source())
  78. }
  79. // TestSyncedUserFunding_BillingSource 测试同步用户计费来源标识
  80. func TestSyncedUserFunding_BillingSource(t *testing.T) {
  81. f := NewSyncedUserFunding(100, 200, "req-src", 500000)
  82. assert.Equal(t, BillingSourceSyncedWallet, f.Source())
  83. assert.Equal(t, "synced_wallet", f.Source())
  84. }
  85. // TestNeedsRefundLocked_WalletAndSynced 对比钱包和同步用户的行为
  86. func TestNeedsRefundLocked_WalletAndSynced(t *testing.T) {
  87. // 钱包用户:tokenConsumed > 0 才需要退
  88. walletSession := &BillingSession{
  89. funding: &WalletFunding{userId: 1, consumed: 100},
  90. tokenConsumed: 100,
  91. }
  92. assert.True(t, walletSession.needsRefundLocked())
  93. walletSessionZero := &BillingSession{
  94. funding: &WalletFunding{userId: 1, consumed: 100},
  95. tokenConsumed: 0,
  96. }
  97. assert.False(t, walletSessionZero.needsRefundLocked())
  98. // 同步用户:consumed > 0 就需要退(即使 tokenConsumed=0)
  99. setupBillingSyncDB(t)
  100. settings := system_setting.GetRegionSyncSettings()
  101. origThreshold := settings.MinBalanceThreshold
  102. defer func() { settings.MinBalanceThreshold = origThreshold }()
  103. settings.MinBalanceThreshold = 100000
  104. syncF := NewSyncedUserFunding(100, 200, "req-compare", 500000)
  105. require.NoError(t, syncF.PreConsume(5000))
  106. syncSession := &BillingSession{
  107. funding: syncF,
  108. tokenConsumed: 0, // tokenConsumed=0 但 consumed > 0
  109. }
  110. assert.True(t, syncSession.needsRefundLocked())
  111. }