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.
 
 
 

539 lines
15 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "testing"
  5. "time"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/glebarez/sqlite"
  8. "github.com/stretchr/testify/assert"
  9. "github.com/stretchr/testify/require"
  10. "gorm.io/gorm"
  11. )
  12. func setupRegionSyncDB(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 := DB
  19. origLogDB := LOG_DB
  20. DB = db
  21. LOG_DB = db
  22. common.UsingSQLite = true
  23. common.RedisEnabled = false
  24. require.NoError(t, db.AutoMigrate(&User{}))
  25. t.Cleanup(func() {
  26. DB = origDB
  27. LOG_DB = origLogDB
  28. sqlDB.Close()
  29. })
  30. return db
  31. }
  32. func TestUser_IsSyncedUser_Method(t *testing.T) {
  33. tests := []struct {
  34. name string
  35. user User
  36. want bool
  37. }{
  38. {"synced valid", User{Source: "synced", Id: 100}, true},
  39. {"synced boundary", User{Source: "synced", Id: common.ForeignUserIDStart - 1}, true},
  40. {"synced id 0", User{Source: "synced", Id: 0}, false},
  41. {"synced too large", User{Source: "synced", Id: common.ForeignUserIDStart}, false},
  42. {"local user", User{Source: "local", Id: 100}, false},
  43. }
  44. for _, tt := range tests {
  45. t.Run(tt.name, func(t *testing.T) {
  46. assert.Equal(t, tt.want, tt.user.IsSyncedUser())
  47. })
  48. }
  49. }
  50. func TestUser_IsLocalUser_Method(t *testing.T) {
  51. tests := []struct {
  52. name string
  53. user User
  54. want bool
  55. }{
  56. {"local source", User{Source: "local", Id: common.ForeignUserIDStart}, true},
  57. {"foreign id", User{Source: "", Id: common.ForeignUserIDStart + 1}, true},
  58. {"synced source", User{Source: "synced", Id: 100}, false},
  59. }
  60. for _, tt := range tests {
  61. t.Run(tt.name, func(t *testing.T) {
  62. assert.Equal(t, tt.want, tt.user.IsLocalUser())
  63. })
  64. }
  65. }
  66. func TestUser_SyncFields_CRUD(t *testing.T) {
  67. db := setupRegionSyncDB(t)
  68. user := User{
  69. Id: 100, Username: "synced_test", Password: "hashed_password",
  70. Source: common.UserSourceSynced, RemoteUserId: 100,
  71. SyncedQuota: 500000, LastSyncAt: 1700000000,
  72. }
  73. require.NoError(t, db.Create(&user).Error)
  74. var loaded User
  75. require.NoError(t, db.First(&loaded, 100).Error)
  76. assert.Equal(t, common.UserSourceSynced, loaded.Source)
  77. assert.Equal(t, 100, loaded.RemoteUserId)
  78. assert.Equal(t, 500000, loaded.SyncedQuota)
  79. assert.Equal(t, int64(1700000000), loaded.LastSyncAt)
  80. }
  81. func TestUser_DefaultSourceIsLocal(t *testing.T) {
  82. db := setupRegionSyncDB(t)
  83. user := User{
  84. Id: 101, Username: "default_source_test", Password: "hashed_password",
  85. }
  86. require.NoError(t, db.Create(&user).Error)
  87. var loaded User
  88. require.NoError(t, db.Where("username = ?", "default_source_test").First(&loaded).Error)
  89. assert.Equal(t, common.UserSourceLocal, loaded.Source)
  90. }
  91. // PendingSyncRecord tests
  92. func setupPendingSyncDB(t *testing.T) *gorm.DB {
  93. t.Helper()
  94. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  95. require.NoError(t, err)
  96. sqlDB, _ := db.DB()
  97. sqlDB.SetMaxOpenConns(1)
  98. origDB := DB
  99. origLogDB := LOG_DB
  100. DB = db
  101. LOG_DB = db
  102. common.UsingSQLite = true
  103. common.RedisEnabled = false
  104. require.NoError(t, db.AutoMigrate(&PendingSyncRecord{}))
  105. t.Cleanup(func() {
  106. DB = origDB
  107. LOG_DB = origLogDB
  108. sqlDB.Close()
  109. })
  110. return db
  111. }
  112. func TestPendingSyncRecord_Create(t *testing.T) {
  113. db := setupPendingSyncDB(t)
  114. userId := 100
  115. remoteUserId := 200
  116. requestId := "req-test-001"
  117. quota := 1000
  118. estimatedQuota := 500
  119. err := CreatePendingSyncRecord(userId, remoteUserId, requestId, quota, estimatedQuota)
  120. require.NoError(t, err)
  121. var record PendingSyncRecord
  122. require.NoError(t, db.Where("request_id = ?", requestId).First(&record).Error)
  123. assert.Equal(t, userId, record.UserId)
  124. assert.Equal(t, remoteUserId, record.RemoteUserId)
  125. assert.Equal(t, requestId, record.RequestId)
  126. assert.Equal(t, quota, record.Quota)
  127. assert.Equal(t, estimatedQuota, record.EstimatedQuota)
  128. assert.Equal(t, PendingSyncStatusPending, record.Status)
  129. assert.NotZero(t, record.CreatedAt)
  130. assert.Zero(t, record.SyncedAt)
  131. assert.Zero(t, record.RetryCount)
  132. }
  133. func TestPendingSyncRecord_RequestIdUnique(t *testing.T) {
  134. db := setupPendingSyncDB(t)
  135. requestId := "req-unique-001"
  136. // 创建第一条记录
  137. err := CreatePendingSyncRecord(100, 200, requestId, 1000, 500)
  138. require.NoError(t, err)
  139. // 尝试用相同的 requestId 创建第二条记录
  140. err = CreatePendingSyncRecord(101, 201, requestId, 2000, 600)
  141. assert.Error(t, err) // 应该失败,因为 requestId 唯一
  142. // 验证只有一条记录
  143. var count int64
  144. db.Model(&PendingSyncRecord{}).Where("request_id = ?", requestId).Count(&count)
  145. assert.Equal(t, int64(1), count)
  146. }
  147. func TestGetPendingSyncQuota(t *testing.T) {
  148. setupPendingSyncDB(t)
  149. userId := 100
  150. // 创建多条记录
  151. require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-001", 1000, 500))
  152. require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-002", 2000, 600))
  153. require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-003", 500, 200))
  154. // 创建一个 failed 状态的记录,应该被计入
  155. require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-004", 300, 100))
  156. DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-004").Update("status", PendingSyncStatusFailed)
  157. // 创建一个 synced 状态的记录,不应该被计入
  158. require.NoError(t, CreatePendingSyncRecord(userId, 200, "req-005", 9999, 0))
  159. DB.Model(&PendingSyncRecord{}).Where("request_id = ?", "req-005").Update("status", PendingSyncStatusSynced)
  160. total := GetPendingSyncQuota(userId)
  161. assert.Equal(t, 3800, total) // 1000 + 2000 + 500 + 300 = 3800
  162. }
  163. func TestGetPendingSyncQuota_Empty(t *testing.T) {
  164. setupPendingSyncDB(t)
  165. // 用户没有待同步记录
  166. total := GetPendingSyncQuota(99999)
  167. assert.Equal(t, 0, total)
  168. }
  169. // QuotaSyncLog tests
  170. func TestQuotaSyncLog_Create(t *testing.T) {
  171. db := setupRegionSyncDB(t)
  172. require.NoError(t, db.AutoMigrate(&QuotaSyncLog{}))
  173. log := QuotaSyncLog{
  174. UserId: 100,
  175. RemoteUserId: 100,
  176. SyncType: SyncTypeUserCreate,
  177. BeforeQuota: 0,
  178. AfterQuota: 500000,
  179. ChangeAmount: 500000,
  180. MasterQuota: 500000,
  181. PendingQuota: 0,
  182. Direction: SyncDirectionCnToOv,
  183. Status: SyncStatusSuccess,
  184. CreatedAt: 1700000000,
  185. SyncedAt: 1700000000,
  186. }
  187. require.NoError(t, db.Create(&log).Error)
  188. var loaded QuotaSyncLog
  189. require.NoError(t, db.First(&loaded, log.Id).Error)
  190. assert.Equal(t, SyncTypeUserCreate, loaded.SyncType)
  191. assert.Equal(t, 500000, loaded.ChangeAmount)
  192. assert.Equal(t, SyncDirectionCnToOv, loaded.Direction)
  193. }
  194. func TestQuotaSyncLog_QueryByUser(t *testing.T) {
  195. db := setupRegionSyncDB(t)
  196. require.NoError(t, db.AutoMigrate(&QuotaSyncLog{}))
  197. for i := 0; i < 3; i++ {
  198. log := QuotaSyncLog{
  199. UserId: 100, RemoteUserId: 100,
  200. SyncType: SyncTypeBatchSync, Direction: SyncDirectionOvToCn,
  201. Status: SyncStatusSuccess, ChangeAmount: 100 * (i + 1),
  202. CreatedAt: 1700000000 + int64(i),
  203. }
  204. require.NoError(t, db.Create(&log).Error)
  205. }
  206. var logs []QuotaSyncLog
  207. db.Where("user_id = ?", 100).Order("created_at desc").Find(&logs)
  208. assert.Len(t, logs, 3)
  209. }
  210. // ---------------------------------------------------------------------------
  211. // 阶段 1:Model 层缺失测试
  212. // ---------------------------------------------------------------------------
  213. func TestGetPendingRecordsForSync_Basic(t *testing.T) {
  214. db := setupPendingSyncDB(t)
  215. now := time.Now().Unix()
  216. require.NoError(t, db.Create(&PendingSyncRecord{
  217. UserId: 100, RemoteUserId: 200, RequestId: "req-rs-001",
  218. Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now - 100,
  219. }).Error)
  220. require.NoError(t, db.Create(&PendingSyncRecord{
  221. UserId: 101, RemoteUserId: 201, RequestId: "req-rs-002",
  222. Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now - 50,
  223. }).Error)
  224. require.NoError(t, db.Create(&PendingSyncRecord{
  225. UserId: 102, RemoteUserId: 202, RequestId: "req-rs-003",
  226. Quota: 300, Status: PendingSyncStatusSynced, CreatedAt: now,
  227. }).Error)
  228. records := GetPendingRecordsForSync(10, 3)
  229. assert.Len(t, records, 2)
  230. // 按 created_at 升序
  231. assert.Equal(t, "req-rs-001", records[0].RequestId)
  232. assert.Equal(t, "req-rs-002", records[1].RequestId)
  233. }
  234. func TestGetPendingRecordsForSync_MaxRetry(t *testing.T) {
  235. db := setupPendingSyncDB(t)
  236. now := time.Now().Unix()
  237. require.NoError(t, db.Create(&PendingSyncRecord{
  238. UserId: 100, RemoteUserId: 200, RequestId: "req-retry-001",
  239. Quota: 100, Status: PendingSyncStatusFailed, RetryCount: 3, CreatedAt: now,
  240. }).Error)
  241. require.NoError(t, db.Create(&PendingSyncRecord{
  242. UserId: 101, RemoteUserId: 201, RequestId: "req-retry-002",
  243. Quota: 200, Status: PendingSyncStatusPending, RetryCount: 2, CreatedAt: now,
  244. }).Error)
  245. records := GetPendingRecordsForSync(10, 3)
  246. assert.Len(t, records, 1)
  247. assert.Equal(t, "req-retry-002", records[0].RequestId)
  248. }
  249. func TestGetPendingRecordsForSync_Limit(t *testing.T) {
  250. db := setupPendingSyncDB(t)
  251. now := time.Now().Unix()
  252. for i := 0; i < 5; i++ {
  253. require.NoError(t, db.Create(&PendingSyncRecord{
  254. UserId: 100, RemoteUserId: 200, RequestId: fmt.Sprintf("req-lim-%d", i),
  255. Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now + int64(i),
  256. }).Error)
  257. }
  258. records := GetPendingRecordsForSync(2, 3)
  259. assert.Len(t, records, 2)
  260. }
  261. func TestMarkRecordSynced(t *testing.T) {
  262. db := setupPendingSyncDB(t)
  263. require.NoError(t, db.Create(&PendingSyncRecord{
  264. UserId: 100, RemoteUserId: 200, RequestId: "req-sync-001",
  265. Quota: 100, Status: PendingSyncStatusPending, CreatedAt: time.Now().Unix(),
  266. }).Error)
  267. var record PendingSyncRecord
  268. require.NoError(t, db.Where("request_id = ?", "req-sync-001").First(&record).Error)
  269. require.NoError(t, MarkRecordSynced(record.Id))
  270. var updated PendingSyncRecord
  271. require.NoError(t, db.First(&updated, record.Id).Error)
  272. assert.Equal(t, PendingSyncStatusSynced, updated.Status)
  273. assert.NotZero(t, updated.SyncedAt)
  274. }
  275. func TestMarkRecordFailed(t *testing.T) {
  276. db := setupPendingSyncDB(t)
  277. require.NoError(t, db.Create(&PendingSyncRecord{
  278. UserId: 100, RemoteUserId: 200, RequestId: "req-fail-001",
  279. Quota: 100, Status: PendingSyncStatusPending, RetryCount: 0, CreatedAt: time.Now().Unix(),
  280. }).Error)
  281. var record PendingSyncRecord
  282. require.NoError(t, db.Where("request_id = ?", "req-fail-001").First(&record).Error)
  283. require.NoError(t, MarkRecordFailed(record.Id, "connection refused"))
  284. var updated PendingSyncRecord
  285. require.NoError(t, db.First(&updated, record.Id).Error)
  286. assert.Equal(t, PendingSyncStatusFailed, updated.Status)
  287. assert.Equal(t, 1, updated.RetryCount)
  288. assert.Equal(t, "connection refused", updated.ErrorMsg)
  289. }
  290. func TestCleanupOldSyncedRecords(t *testing.T) {
  291. db := setupPendingSyncDB(t)
  292. now := time.Now().Unix()
  293. // 40 天前已同步的记录
  294. require.NoError(t, db.Create(&PendingSyncRecord{
  295. UserId: 100, RemoteUserId: 200, RequestId: "req-old-001",
  296. Quota: 100, Status: PendingSyncStatusSynced, SyncedAt: now - 40*86400, CreatedAt: now - 45*86400,
  297. }).Error)
  298. // 5 天前已同步的记录(不应被清理)
  299. require.NoError(t, db.Create(&PendingSyncRecord{
  300. UserId: 100, RemoteUserId: 200, RequestId: "req-old-002",
  301. Quota: 200, Status: PendingSyncStatusSynced, SyncedAt: now - 5*86400, CreatedAt: now - 10*86400,
  302. }).Error)
  303. // pending 状态的记录(不应被清理)
  304. require.NoError(t, db.Create(&PendingSyncRecord{
  305. UserId: 100, RemoteUserId: 200, RequestId: "req-old-003",
  306. Quota: 300, Status: PendingSyncStatusPending, CreatedAt: now - 50*86400,
  307. }).Error)
  308. deleted := CleanupOldSyncedRecords(30)
  309. assert.Equal(t, int64(1), deleted)
  310. var count int64
  311. db.Model(&PendingSyncRecord{}).Count(&count)
  312. assert.Equal(t, int64(2), count)
  313. }
  314. func TestCleanupOldSyncedRecords_Empty(t *testing.T) {
  315. setupPendingSyncDB(t)
  316. deleted := CleanupOldSyncedRecords(30)
  317. assert.Equal(t, int64(0), deleted)
  318. }
  319. func TestGetPendingQuotaByUser(t *testing.T) {
  320. db := setupPendingSyncDB(t)
  321. now := time.Now().Unix()
  322. require.NoError(t, db.Create(&PendingSyncRecord{
  323. UserId: 100, RemoteUserId: 200, RequestId: "req-pq-001",
  324. Quota: 100, Status: PendingSyncStatusPending, CreatedAt: now,
  325. }).Error)
  326. require.NoError(t, db.Create(&PendingSyncRecord{
  327. UserId: 100, RemoteUserId: 200, RequestId: "req-pq-002",
  328. Quota: 200, Status: PendingSyncStatusFailed, CreatedAt: now,
  329. }).Error)
  330. require.NoError(t, db.Create(&PendingSyncRecord{
  331. UserId: 101, RemoteUserId: 201, RequestId: "req-pq-003",
  332. Quota: 500, Status: PendingSyncStatusPending, CreatedAt: now,
  333. }).Error)
  334. require.NoError(t, db.Create(&PendingSyncRecord{
  335. UserId: 100, RemoteUserId: 200, RequestId: "req-pq-004",
  336. Quota: 9999, Status: PendingSyncStatusSynced, CreatedAt: now,
  337. }).Error)
  338. m := GetPendingQuotaByUser()
  339. assert.Equal(t, 300, m[100]) // 100 + 200, synced 不计入
  340. assert.Equal(t, 500, m[101])
  341. _, exists := m[102]
  342. assert.False(t, exists)
  343. }
  344. // ---------------------------------------------------------------------------
  345. // QuotaUpdateCallback tests
  346. // ---------------------------------------------------------------------------
  347. func TestSetQuotaUpdateCallback(t *testing.T) {
  348. // 测试设置回调函数
  349. called := false
  350. var callbackUserId int
  351. var callbackQuota int
  352. callback := func(userId int, quota int) {
  353. called = true
  354. callbackUserId = userId
  355. callbackQuota = quota
  356. }
  357. SetQuotaUpdateCallback(callback)
  358. // 清理:重置为 nil
  359. defer SetQuotaUpdateCallback(nil)
  360. // 验证回调被设置(通过调用测试)
  361. if quotaUpdateCallback != nil {
  362. quotaUpdateCallback(123, 500000)
  363. }
  364. assert.True(t, called)
  365. assert.Equal(t, 123, callbackUserId)
  366. assert.Equal(t, 500000, callbackQuota)
  367. }
  368. func TestIncreaseUserQuota_TriggersCallback(t *testing.T) {
  369. db := setupRegionSyncDB(t)
  370. // 创建测试用户
  371. user := User{
  372. Id: 1000,
  373. Username: "callback_test",
  374. Password: "hashed_password",
  375. Quota: 100000,
  376. }
  377. require.NoError(t, db.Create(&user).Error)
  378. // 设置回调
  379. callbackCalled := false
  380. var receivedUserId int
  381. var receivedQuota int
  382. SetQuotaUpdateCallback(func(userId int, quota int) {
  383. callbackCalled = true
  384. receivedUserId = userId
  385. receivedQuota = quota
  386. })
  387. defer SetQuotaUpdateCallback(nil) // 清理
  388. // 执行增加额度操作
  389. err := IncreaseUserQuota(1000, 50000, true)
  390. require.NoError(t, err)
  391. // 验证回调被触发
  392. assert.True(t, callbackCalled)
  393. assert.Equal(t, 1000, receivedUserId)
  394. assert.Equal(t, 150000, receivedQuota) // 原始 100000 + 增加 50000
  395. // 验证数据库已更新
  396. var updated User
  397. require.NoError(t, db.First(&updated, 1000).Error)
  398. assert.Equal(t, 150000, updated.Quota)
  399. }
  400. func TestIncreaseUserQuota_NoCallback(t *testing.T) {
  401. db := setupRegionSyncDB(t)
  402. // 创建测试用户
  403. user := User{
  404. Id: 1001,
  405. Username: "no_callback_test",
  406. Password: "hashed_password",
  407. Quota: 100000,
  408. }
  409. require.NoError(t, db.Create(&user).Error)
  410. // 确保没有回调
  411. SetQuotaUpdateCallback(nil)
  412. // 执行增加额度操作(不应 panic)
  413. err := IncreaseUserQuota(1001, 50000, true)
  414. require.NoError(t, err)
  415. // 验证数据库已更新
  416. var updated User
  417. require.NoError(t, db.First(&updated, 1001).Error)
  418. assert.Equal(t, 150000, updated.Quota)
  419. }
  420. func TestIncreaseUserQuota_CallbackError(t *testing.T) {
  421. db := setupRegionSyncDB(t)
  422. // 创建测试用户
  423. user := User{
  424. Id: 1002,
  425. Username: "callback_error_test",
  426. Password: "hashed_password",
  427. Quota: 100000,
  428. }
  429. require.NoError(t, db.Create(&user).Error)
  430. // 设置一个会 panic 的回调,验证不影响主流程
  431. SetQuotaUpdateCallback(func(userId int, quota int) {
  432. panic("callback error")
  433. })
  434. defer SetQuotaUpdateCallback(nil)
  435. // 执行增加额度操作(回调 panic 不应影响数据库更新)
  436. err := IncreaseUserQuota(1002, 50000, true)
  437. require.NoError(t, err)
  438. // 验证数据库已更新(即使回调失败)
  439. var updated User
  440. require.NoError(t, db.First(&updated, 1002).Error)
  441. assert.Equal(t, 150000, updated.Quota)
  442. }