|
- package model
-
- import (
- "fmt"
- "testing"
- "time"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupRegionSyncDB(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 := DB
- origLogDB := LOG_DB
- DB = db
- LOG_DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&User{}))
-
- t.Cleanup(func() {
- DB = origDB
- LOG_DB = origLogDB
- sqlDB.Close()
- })
- return db
- }
-
- func TestUser_IsSyncedUser_Method(t *testing.T) {
- tests := []struct {
- name string
- user User
- want bool
- }{
- {"synced with normal id", User{Source: "synced", Id: 100}, true},
- {"synced with large id", User{Source: "synced", Id: 99999999}, true},
- {"synced with id 0", User{Source: "synced", Id: 0}, true},
- {"local user", User{Source: "local", Id: 100}, false},
- {"empty source", User{Source: "", Id: 100}, false},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- assert.Equal(t, tt.want, tt.user.IsSyncedUser())
- })
- }
- }
-
- func TestUser_IsLocalUser_Method(t *testing.T) {
- tests := []struct {
- name string
- user User
- want bool
- }{
- {"local source", User{Source: "local", Id: 100}, true},
- {"empty source is local", User{Source: "", Id: 100}, true},
- {"synced source is not local", User{Source: "synced", Id: 100}, false},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- assert.Equal(t, tt.want, tt.user.IsLocalUser())
- })
- }
- }
-
- func TestUser_SyncFields_CRUD(t *testing.T) {
- db := setupRegionSyncDB(t)
-
- user := User{
- Id: 100, Username: "synced_test", Password: "hashed_password",
- Source: common.UserSourceSynced, RemoteUserId: 100,
- SyncedQuota: 500000, LastSyncAt: 1700000000,
- }
- require.NoError(t, db.Create(&user).Error)
-
- var loaded User
- require.NoError(t, db.First(&loaded, 100).Error)
- assert.Equal(t, common.UserSourceSynced, loaded.Source)
- assert.Equal(t, 100, loaded.RemoteUserId)
- assert.Equal(t, 500000, loaded.SyncedQuota)
- assert.Equal(t, int64(1700000000), loaded.LastSyncAt)
- }
-
- func TestUser_DefaultSourceIsLocal(t *testing.T) {
- db := setupRegionSyncDB(t)
-
- user := User{
- Id: 101, Username: "default_source_test", Password: "hashed_password",
- }
- require.NoError(t, db.Create(&user).Error)
-
- var loaded User
- require.NoError(t, db.Where("username = ?", "default_source_test").First(&loaded).Error)
- assert.Equal(t, common.UserSourceLocal, loaded.Source)
- }
-
- // PendingSyncRecord tests
-
- func setupPendingSyncDB(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 := DB
- origLogDB := LOG_DB
- DB = db
- LOG_DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&PendingSyncRecord{}))
-
- t.Cleanup(func() {
- DB = origDB
- LOG_DB = origLogDB
- sqlDB.Close()
- })
- return db
- }
-
- func TestPendingSyncRecord_Create(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- userId := 100
- remoteUserId := 200
- requestId := "req-test-001"
- quota := 1000
- estimatedQuota := 500
-
- err := CreatePendingSyncRecord(userId, remoteUserId, requestId, quota, estimatedQuota)
- require.NoError(t, err)
-
- var record PendingSyncRecord
- require.NoError(t, db.Where("request_id = ?", requestId).First(&record).Error)
-
- assert.Equal(t, userId, record.UserId)
- assert.Equal(t, remoteUserId, record.RemoteUserId)
- assert.Equal(t, requestId, record.RequestId)
- assert.Equal(t, quota, record.Quota)
- assert.Equal(t, estimatedQuota, record.EstimatedQuota)
- assert.Equal(t, PendingSyncStatusPending, record.Status)
- assert.NotZero(t, record.CreatedAt)
- assert.Zero(t, record.SyncedAt)
- assert.Zero(t, record.RetryCount)
- }
-
- func TestPendingSyncRecord_RequestIdUnique(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- requestId := "req-unique-001"
-
- // 创建第一条记录
- err := CreatePendingSyncRecord(100, 200, requestId, 1000, 500)
- require.NoError(t, err)
-
- // 尝试用相同的 requestId 创建第二条记录
- err = CreatePendingSyncRecord(101, 201, requestId, 2000, 600)
- assert.Error(t, err) // 应该失败,因为 requestId 唯一
-
- // 验证只有一条记录
- var count int64
- db.Model(&PendingSyncRecord{}).Where("request_id = ?", requestId).Count(&count)
- assert.Equal(t, int64(1), count)
- }
-
- func TestGetPendingSyncQuota(t *testing.T) {
- setupPendingSyncDB(t)
- userId := 100
-
- // 创建多条记录
- require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-001", 1000, 500))
- require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-002", 2000, 600))
- require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-003", 500, 200))
-
- // 创建一个 failed 状态的记录,应该被计入
- require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-004", 300, 100))
- DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-004").Update("status", PendingSyncStatusFailed)
-
- // 创建一个 settled 状态的记录,应该被计入
- require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-005", 400, 200))
- DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-005").Update("status", PendingSyncStatusSettled)
-
- // 创建一个 synced 状态的记录,不应该被计入
- require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-006", 9999, 0))
- DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-006").Update("status", PendingSyncStatusSynced)
-
- total := GetPendingSyncQuota(userId)
- assert.Equal(t, 4200, total) // 1000 + 2000 + 500 + 300 + 400 = 4200
- }
-
- func TestGetPendingSyncQuota_Empty(t *testing.T) {
- setupPendingSyncDB(t)
-
- // 用户没有待同步记录
- total := GetPendingSyncQuota(99999)
- assert.Equal(t, 0, total)
- }
-
- // QuotaSyncLog tests
-
- func TestQuotaSyncLog_Create(t *testing.T) {
- db := setupRegionSyncDB(t)
- require.NoError(t, db.AutoMigrate(&QuotaSyncLog{}))
-
- log := QuotaSyncLog{
- UserId: 100,
- RemoteUserId: 100,
- SyncType: SyncTypeUserCreate,
- BeforeQuota: 0,
- AfterQuota: 500000,
- ChangeAmount: 500000,
- MasterQuota: 500000,
- PendingQuota: 0,
- Direction: SyncDirectionCnToOv,
- Status: SyncStatusSuccess,
- CreatedAt: 1700000000,
- SyncedAt: 1700000000,
- }
- require.NoError(t, db.Create(&log).Error)
-
- var loaded QuotaSyncLog
- require.NoError(t, db.First(&loaded, log.Id).Error)
- assert.Equal(t, SyncTypeUserCreate, loaded.SyncType)
- assert.Equal(t, 500000, loaded.ChangeAmount)
- assert.Equal(t, SyncDirectionCnToOv, loaded.Direction)
- }
-
- func TestQuotaSyncLog_QueryByUser(t *testing.T) {
- db := setupRegionSyncDB(t)
- require.NoError(t, db.AutoMigrate(&QuotaSyncLog{}))
-
- for i := 0; i < 3; i++ {
- log := QuotaSyncLog{
- UserId: 100, RemoteUserId: 100,
- SyncType: SyncTypeBatchSync, Direction: SyncDirectionOvToCn,
- Status: SyncStatusSuccess, ChangeAmount: 100 * (i + 1),
- CreatedAt: 1700000000 + int64(i),
- }
- require.NoError(t, db.Create(&log).Error)
- }
-
- var logs []QuotaSyncLog
- db.Where("user_id = ?", 100).Order("created_at desc").Find(&logs)
- assert.Len(t, logs, 3)
- }
-
- // ---------------------------------------------------------------------------
- // 阶段 1:Model 层缺失测试
- // ---------------------------------------------------------------------------
-
- func TestGetPendingRecordsForSync_Basic(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- // settled 记录会被同步
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-rs-001",
- Quota: 100, Status: PendingSyncStatusSettled, CreatedAt: now - 100,
- }).Error)
- // failed 记录会被重试
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 101, RemoteUserId: 201, RequestId: "req-rs-002",
- Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now - 50,
- }).Error)
- // synced 记录不会被查到
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 102, RemoteUserId: 202, RequestId: "req-rs-003",
- Quota: 300, Status: PendingSyncStatusSynced, CreatedAt: now,
- }).Error)
- // pending 记录不会被同步(等待 Settle)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 103, RemoteUserId: 203, RequestId: "req-rs-004",
- Quota: 400, Status: PendingSyncStatusPending, CreatedAt: now + 1,
- }).Error)
-
- records := GetPendingRecordsForSync(10, 3)
- assert.Len(t, records, 2) // 只查 settled + failed
- // 按 created_at 升序
- assert.Equal(t, "req-rs-001", records[0].RequestId)
- assert.Equal(t, "req-rs-002", records[1].RequestId)
- }
-
- func TestGetPendingRecordsForSync_MaxRetry(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-retry-001",
- Quota: 100, Status: PendingSyncStatusFailed, RetryCount: 3, CreatedAt: now,
- }).Error)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 101, RemoteUserId: 201, RequestId: "req-retry-002",
- Quota: 200, Status: PendingSyncStatusSettled, RetryCount: 2, CreatedAt: now,
- }).Error)
-
- records := GetPendingRecordsForSync(10, 3)
- assert.Len(t, records, 1)
- assert.Equal(t, "req-retry-002", records[0].RequestId)
- }
-
- func TestGetPendingRecordsForSync_Limit(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- for i := 0; i < 5; i++ {
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: fmt.Sprintf("req-lim-%d", i),
- Quota: 100, Status: PendingSyncStatusSettled, CreatedAt: now + int64(i),
- }).Error)
- }
-
- records := GetPendingRecordsForSync(2, 3)
- assert.Len(t, records, 2)
- }
-
- func TestMarkRecordSynced(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-sync-001",
- Quota: 100, Status: PendingSyncStatusPending, CreatedAt: time.Now().Unix(),
- }).Error)
-
- var record PendingSyncRecord
- require.NoError(t, db.Where("request_id = ?", "req-sync-001").First(&record).Error)
-
- require.NoError(t, MarkRecordSynced(record.Id))
-
- var updated PendingSyncRecord
- require.NoError(t, db.First(&updated, record.Id).Error)
- assert.Equal(t, PendingSyncStatusSynced, updated.Status)
- assert.NotZero(t, updated.SyncedAt)
- }
-
- func TestMarkRecordFailed(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-fail-001",
- Quota: 100, Status: PendingSyncStatusPending, RetryCount: 0, CreatedAt: time.Now().Unix(),
- }).Error)
-
- var record PendingSyncRecord
- require.NoError(t, db.Where("request_id = ?", "req-fail-001").First(&record).Error)
-
- require.NoError(t, MarkRecordFailed(record.Id, "connection refused"))
-
- var updated PendingSyncRecord
- require.NoError(t, db.First(&updated, record.Id).Error)
- assert.Equal(t, PendingSyncStatusFailed, updated.Status)
- assert.Equal(t, 1, updated.RetryCount)
- assert.Equal(t, "connection refused", updated.ErrorMsg)
- }
-
- func TestCleanupOldSyncedRecords(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- // 40 天前已同步的记录
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-old-001",
- Quota: 100, Status: PendingSyncStatusSynced, SyncedAt: now - 40*86400, CreatedAt: now - 45*86400,
- }).Error)
- // 5 天前已同步的记录(不应被清理)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-old-002",
- Quota: 200, Status: PendingSyncStatusSynced, SyncedAt: now - 5*86400, CreatedAt: now - 10*86400,
- }).Error)
- // pending 状态的记录(不应被清理)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-old-003",
- Quota: 300, Status: PendingSyncStatusPending, CreatedAt: now - 50*86400,
- }).Error)
-
- deleted := CleanupOldSyncedRecords(30)
- assert.Equal(t, int64(1), deleted)
-
- var count int64
- db.Model(&PendingSyncRecord{}).Count(&count)
- assert.Equal(t, int64(2), count)
- }
-
- func TestCleanupOldSyncedRecords_Empty(t *testing.T) {
- setupPendingSyncDB(t)
-
- deleted := CleanupOldSyncedRecords(30)
- assert.Equal(t, int64(0), deleted)
- }
-
- func TestGetPendingQuotaByUser(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-pq-001",
- Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now,
- }).Error)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-pq-002",
- Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now,
- }).Error)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 101, RemoteUserId: 201, RequestId: "req-pq-003",
- Quota: 500, Status: PendingSyncStatusPending, CreatedAt: now,
- }).Error)
- // settled 记录应计入
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-pq-004",
- Quota: 150, Status: PendingSyncStatusSettled, CreatedAt: now,
- }).Error)
- // synced 记录不计入
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-pq-005",
- Quota: 9999, Status: PendingSyncStatusSynced, CreatedAt: now,
- }).Error)
-
- m := GetPendingQuotaByUser()
- assert.Equal(t, 450, m[100]) // 100 + 200 + 150 (pending + failed + settled), synced 不计入
- assert.Equal(t, 500, m[101])
- _, exists := m[102]
- assert.False(t, exists)
- }
-
-
- // ---------------------------------------------------------------------------
- // Settled status tests
- // ---------------------------------------------------------------------------
-
- func TestSettlePendingSyncRecord(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-settle-1",
- Quota: 1000, EstimatedQuota: 1000, Status: PendingSyncStatusPending, CreatedAt: now,
- }).Error)
-
- found, err := SettlePendingSyncRecord("req-settle-1", 700)
- require.NoError(t, err)
- assert.True(t, found)
-
- var record PendingSyncRecord
- require.NoError(t, db.Where("request_id = ?", "req-settle-1").First(&record).Error)
- assert.Equal(t, PendingSyncStatusSettled, record.Status)
- assert.Equal(t, 700, record.Quota)
- }
-
- func TestSettlePendingSyncRecord_NotFound(t *testing.T) {
- setupPendingSyncDB(t)
-
- found, err := SettlePendingSyncRecord("nonexistent", 100)
- require.NoError(t, err)
- assert.False(t, found)
- }
-
- func TestCreateSettledSyncRecord(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- err := CreateSettledSyncRecord(100, 200, "req-settled-direct", 500)
- require.NoError(t, err)
-
- var record PendingSyncRecord
- require.NoError(t, db.Where("request_id = ?", "req-settled-direct").First(&record).Error)
- assert.Equal(t, PendingSyncStatusSettled, record.Status)
- assert.Equal(t, 500, record.Quota)
- }
-
- func TestArchiveStalePendingRecords(t *testing.T) {
- db := setupPendingSyncDB(t)
-
- now := time.Now().Unix()
- // 超时的 pending 记录(应被归档)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-stale-1",
- Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now - 600,
- }).Error)
- // 未超时的 pending 记录(不应被归档)
- require.NoError(t, db.Create(&PendingSyncRecord{
- UserId: 100, RemoteUserId: 200, RequestId: "req-stale-2",
- Quota: 200, Status: PendingSyncStatusPending, CreatedAt: now - 100,
- }).Error)
-
- archived := ArchiveStalePendingRecords(300) // 5分钟超时
- assert.Equal(t, int64(1), archived)
-
- var record PendingSyncRecord
- require.NoError(t, db.Where("request_id = ?", "req-stale-1").First(&record).Error)
- assert.Equal(t, PendingSyncStatusArchived, record.Status)
- }
-
- // ---------------------------------------------------------------------------
- // QuotaUpdateCallback tests
- // ---------------------------------------------------------------------------
-
- func TestSetQuotaUpdateCallback(t *testing.T) {
- // 测试设置回调函数
- called := false
- var callbackUserId int
- var callbackQuota int
-
- callback := func(userId int, quota int) {
- called = true
- callbackUserId = userId
- callbackQuota = quota
- }
-
- SetQuotaUpdateCallback(callback)
-
- // 清理:重置为 nil
- defer SetQuotaUpdateCallback(nil)
-
- // 验证回调被设置(通过调用测试)
- if quotaUpdateCallback != nil {
- quotaUpdateCallback(123, 500000)
- }
-
- assert.True(t, called)
- assert.Equal(t, 123, callbackUserId)
- assert.Equal(t, 500000, callbackQuota)
- }
-
- func TestIncreaseUserQuota_TriggersCallback(t *testing.T) {
- db := setupRegionSyncDB(t)
-
- // 创建测试用户
- user := User{
- Id: 1000,
- Username: "callback_test",
- Password: "hashed_password",
- Quota: 100000,
- }
- require.NoError(t, db.Create(&user).Error)
-
- // 设置回调
- callbackCalled := false
- var receivedUserId int
- var receivedQuota int
-
- SetQuotaUpdateCallback(func(userId int, quota int) {
- callbackCalled = true
- receivedUserId = userId
- receivedQuota = quota
- })
- defer SetQuotaUpdateCallback(nil) // 清理
-
- // 执行增加额度操作
- err := IncreaseUserQuota(1000, 50000, true)
- require.NoError(t, err)
-
- // 验证回调被触发
- assert.True(t, callbackCalled)
- assert.Equal(t, 1000, receivedUserId)
- assert.Equal(t, 150000, receivedQuota) // 原始 100000 + 增加 50000
-
- // 验证数据库已更新
- var updated User
- require.NoError(t, db.First(&updated, 1000).Error)
- assert.Equal(t, 150000, updated.Quota)
- }
-
- func TestIncreaseUserQuota_NoCallback(t *testing.T) {
- db := setupRegionSyncDB(t)
-
- // 创建测试用户
- user := User{
- Id: 1001,
- Username: "no_callback_test",
- Password: "hashed_password",
- Quota: 100000,
- }
- require.NoError(t, db.Create(&user).Error)
-
- // 确保没有回调
- SetQuotaUpdateCallback(nil)
-
- // 执行增加额度操作(不应 panic)
- err := IncreaseUserQuota(1001, 50000, true)
- require.NoError(t, err)
-
- // 验证数据库已更新
- var updated User
- require.NoError(t, db.First(&updated, 1001).Error)
- assert.Equal(t, 150000, updated.Quota)
- }
-
- func TestIncreaseUserQuota_CallbackError(t *testing.T) {
- db := setupRegionSyncDB(t)
-
- // 创建测试用户
- user := User{
- Id: 1002,
- Username: "callback_error_test",
- Password: "hashed_password",
- Quota: 100000,
- }
- require.NoError(t, db.Create(&user).Error)
-
- // 设置一个会 panic 的回调,验证不影响主流程
- SetQuotaUpdateCallback(func(userId int, quota int) {
- panic("callback error")
- })
- defer SetQuotaUpdateCallback(nil)
-
- // 执行增加额度操作(回调 panic 不应影响数据库更新)
- err := IncreaseUserQuota(1002, 50000, true)
- require.NoError(t, err)
-
- // 验证数据库已更新(即使回调失败)
- var updated User
- require.NoError(t, db.First(&updated, 1002).Error)
- assert.Equal(t, 150000, updated.Quota)
- }
|