|
- package service
-
- import (
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "github.com/QuantumNous/new-api/setting/system_setting"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupBillingSyncDB(t *testing.T) *gorm.DB {
- t.Helper()
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := model.DB
- origLogDB := model.LOG_DB
- model.DB = db
- model.LOG_DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&model.PendingSyncRecord{}, &model.User{}))
-
- t.Cleanup(func() {
- model.DB = origDB
- model.LOG_DB = origLogDB
- sqlDB.Close()
- })
- return db
- }
-
- // TestNeedsRefundLocked_SyncedUser 测试同步用户在 consumed > 0 时需要退款
- func TestNeedsRefundLocked_SyncedUser(t *testing.T) {
- setupBillingSyncDB(t)
-
- settings := system_setting.GetRegionSyncSettings()
- origThreshold := settings.MinBalanceThreshold
- defer func() { settings.MinBalanceThreshold = origThreshold }()
- settings.MinBalanceThreshold = 100000
-
- require.NoError(t, model.DB.Create(&model.User{Id: 100, Username: "test-nr", Source: common.UserSourceSynced, SyncedQuota: 500000}).Error)
- f := NewSyncedUserFunding(100, 200, "req-nr-001", 500000)
- require.NoError(t, f.PreConsume(5000))
-
- session := &BillingSession{
- funding: f,
- }
- assert.True(t, session.needsRefundLocked())
- }
-
- // TestNeedsRefundLocked_SyncedUser_NoConsumed 测试同步用户在 consumed == 0 时不需要退款
- func TestNeedsRefundLocked_SyncedUser_NoConsumed(t *testing.T) {
- setupBillingSyncDB(t)
-
- f := NewSyncedUserFunding(100, 200, "req-nr-002", 500000)
-
- session := &BillingSession{
- funding: f,
- }
- assert.False(t, session.needsRefundLocked())
- }
-
- // TestNeedsRefundLocked_SyncedUser_FundingSettled 测试 fundingSettled=true 后不需要退款
- func TestNeedsRefundLocked_SyncedUser_FundingSettled(t *testing.T) {
- setupBillingSyncDB(t)
-
- settings := system_setting.GetRegionSyncSettings()
- origThreshold := settings.MinBalanceThreshold
- defer func() { settings.MinBalanceThreshold = origThreshold }()
- settings.MinBalanceThreshold = 100000
-
- require.NoError(t, model.DB.Create(&model.User{Id: 100, Username: "test-nr3", Source: common.UserSourceSynced, SyncedQuota: 500000}).Error)
- f := NewSyncedUserFunding(100, 200, "req-nr-003", 500000)
- require.NoError(t, f.PreConsume(5000))
-
- session := &BillingSession{
- funding: f,
- fundingSettled: true, // 资金已提交
- }
- // fundingSettled=true 时不应再退
- assert.False(t, session.needsRefundLocked())
- }
-
- // TestShouldTrust_SyncedUser_SourceNotWalletOrSub 验证同步用户 Source 不是 wallet/subscription,
- // 因此在 shouldTrust 的 switch 中走 default 分支返回 false
- func TestShouldTrust_SyncedUser_SourceNotWalletOrSub(t *testing.T) {
- f := NewSyncedUserFunding(100, 200, "req-trust", 500000)
- assert.Equal(t, "synced_wallet", f.Source())
- assert.NotEqual(t, BillingSourceWallet, f.Source())
- assert.NotEqual(t, BillingSourceSubscription, f.Source())
- }
-
- // TestSyncedUserFunding_BillingSource 测试同步用户计费来源标识
- func TestSyncedUserFunding_BillingSource(t *testing.T) {
- f := NewSyncedUserFunding(100, 200, "req-src", 500000)
- assert.Equal(t, BillingSourceSyncedWallet, f.Source())
- assert.Equal(t, "synced_wallet", f.Source())
- }
-
- // TestNeedsRefundLocked_WalletAndSynced 对比钱包和同步用户的行为
- func TestNeedsRefundLocked_WalletAndSynced(t *testing.T) {
- // 钱包用户:tokenConsumed > 0 才需要退
- walletSession := &BillingSession{
- funding: &WalletFunding{userId: 1, consumed: 100},
- tokenConsumed: 100,
- }
- assert.True(t, walletSession.needsRefundLocked())
-
- walletSessionZero := &BillingSession{
- funding: &WalletFunding{userId: 1, consumed: 100},
- tokenConsumed: 0,
- }
- assert.False(t, walletSessionZero.needsRefundLocked())
-
- // 同步用户:consumed > 0 就需要退(即使 tokenConsumed=0)
- setupBillingSyncDB(t)
- settings := system_setting.GetRegionSyncSettings()
- origThreshold := settings.MinBalanceThreshold
- defer func() { settings.MinBalanceThreshold = origThreshold }()
- settings.MinBalanceThreshold = 100000
-
- require.NoError(t, model.DB.Create(&model.User{Id: 100, Username: "test-compare", Source: common.UserSourceSynced, SyncedQuota: 500000}).Error)
- syncF := NewSyncedUserFunding(100, 200, "req-compare", 500000)
- require.NoError(t, syncF.PreConsume(5000))
-
- syncSession := &BillingSession{
- funding: syncF,
- tokenConsumed: 0, // tokenConsumed=0 但 consumed > 0
- }
- assert.True(t, syncSession.needsRefundLocked())
- }
|