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ů.
 
 
 

467 řádky
13 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "strings"
  8. "testing"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/model"
  11. "github.com/QuantumNous/new-api/service/region_sync"
  12. "github.com/gin-gonic/gin"
  13. "github.com/glebarez/sqlite"
  14. "github.com/stretchr/testify/assert"
  15. "github.com/stretchr/testify/require"
  16. "gorm.io/gorm"
  17. )
  18. func TestReceiveSyncedUserCreate_Success(t *testing.T) {
  19. gin.SetMode(gin.TestMode)
  20. router := gin.New()
  21. // 使用测试专用的认证中间件,绕过系统设置检查
  22. router.Use(testSyncAuthMiddleware())
  23. router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate)
  24. origDB := model.DB
  25. defer func() {
  26. model.DB = origDB
  27. }()
  28. // 设置测试数据库
  29. db := setupSyncTestDB(t)
  30. model.DB = db
  31. req := region_sync.SyncUserRequest{
  32. Username: "testuser",
  33. Email: "test@example.com",
  34. PasswordHash: "hashedpassword",
  35. DisplayName: "Test User",
  36. Quota: 10000,
  37. RemoteUserId: 10000001,
  38. Group: "default",
  39. }
  40. reqBody, _ := json.Marshal(req)
  41. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", bytes.NewBuffer(reqBody))
  42. reqHTTP.Header.Set("Content-Type", "application/json")
  43. w := httptest.NewRecorder()
  44. router.ServeHTTP(w, reqHTTP)
  45. assert.Equal(t, http.StatusOK, w.Code)
  46. var resp map[string]interface{}
  47. json.Unmarshal(w.Body.Bytes(), &resp)
  48. assert.True(t, resp["success"].(bool))
  49. // 验证用户已创建
  50. var user model.User
  51. err := db.Where("username = ?", "testuser").First(&user).Error
  52. require.NoError(t, err)
  53. assert.Equal(t, "testuser", user.Username)
  54. assert.Equal(t, 10000001, user.RemoteUserId)
  55. assert.Equal(t, "synced", user.Source)
  56. }
  57. func TestReceiveQuotaUpdate_Success(t *testing.T) {
  58. gin.SetMode(gin.TestMode)
  59. router := gin.New()
  60. router.Use(testSyncAuthMiddleware())
  61. router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate)
  62. origDB := model.DB
  63. defer func() {
  64. model.DB = origDB
  65. }()
  66. // 设置测试数据库
  67. db := setupSyncTestDB(t)
  68. model.DB = db
  69. // 创建测试用户
  70. user := model.User{
  71. Username: "testuser",
  72. RemoteUserId: 10000001,
  73. Source: "synced",
  74. Quota: 5000,
  75. SyncedQuota: 5000,
  76. AffCode: "aff-qu-1",
  77. }
  78. db.Create(&user)
  79. req := region_sync.UpdateQuotaRequest{
  80. RemoteUserId: 10000001,
  81. Quota: 10000,
  82. }
  83. reqBody, _ := json.Marshal(req)
  84. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", bytes.NewBuffer(reqBody))
  85. reqHTTP.Header.Set("Content-Type", "application/json")
  86. w := httptest.NewRecorder()
  87. router.ServeHTTP(w, reqHTTP)
  88. assert.Equal(t, http.StatusOK, w.Code)
  89. var resp map[string]interface{}
  90. json.Unmarshal(w.Body.Bytes(), &resp)
  91. assert.True(t, resp["success"].(bool))
  92. // 验证余额已更新
  93. var updatedUser model.User
  94. err := db.Where("remote_user_id = ?", 10000001).First(&updatedUser).Error
  95. require.NoError(t, err)
  96. assert.Equal(t, 10000, updatedUser.SyncedQuota)
  97. }
  98. func TestQueryUserQuota_Success(t *testing.T) {
  99. gin.SetMode(gin.TestMode)
  100. router := gin.New()
  101. router.Use(testSyncAuthMiddleware())
  102. router.POST("/api/internal/sync/quota/query", QueryUserQuota)
  103. origDB := model.DB
  104. defer func() {
  105. model.DB = origDB
  106. }()
  107. // 设置测试数据库
  108. db := setupSyncTestDB(t)
  109. model.DB = db
  110. // 创建测试用户
  111. user := model.User{
  112. Username: "testuser",
  113. Quota: 15000,
  114. AffCode: "aff-qq-1",
  115. }
  116. db.Create(&user)
  117. req := region_sync.QueryQuotaRequest{
  118. UserId: user.Id,
  119. }
  120. reqBody, _ := json.Marshal(req)
  121. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/query", bytes.NewBuffer(reqBody))
  122. reqHTTP.Header.Set("Content-Type", "application/json")
  123. w := httptest.NewRecorder()
  124. router.ServeHTTP(w, reqHTTP)
  125. assert.Equal(t, http.StatusOK, w.Code)
  126. var resp region_sync.QueryQuotaResponse
  127. json.Unmarshal(w.Body.Bytes(), &resp)
  128. assert.True(t, resp.Success)
  129. assert.Equal(t, 15000, resp.Quota)
  130. }
  131. func TestBatchDeductQuota_Success(t *testing.T) {
  132. gin.SetMode(gin.TestMode)
  133. router := gin.New()
  134. router.Use(testSyncAuthMiddleware())
  135. router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota)
  136. origDB := model.DB
  137. defer func() {
  138. model.DB = origDB
  139. }()
  140. // 设置测试数据库
  141. db := setupSyncTestDB(t)
  142. model.DB = db
  143. // 创建测试用户
  144. user1 := model.User{Username: "user1", Quota: 1000, AffCode: "aff-bd-1"}
  145. user2 := model.User{Username: "user2", Quota: 500, AffCode: "aff-bd-2"}
  146. db.Create(&user1)
  147. db.Create(&user2)
  148. req := region_sync.BatchDeductRequest{
  149. Records: []region_sync.BatchDeductRecord{
  150. {UserId: user1.Id, RequestId: "req-1", Quota: 100},
  151. {UserId: user2.Id, RequestId: "req-2", Quota: 50},
  152. },
  153. }
  154. reqBody, _ := json.Marshal(req)
  155. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody))
  156. reqHTTP.Header.Set("Content-Type", "application/json")
  157. w := httptest.NewRecorder()
  158. router.ServeHTTP(w, reqHTTP)
  159. assert.Equal(t, http.StatusOK, w.Code)
  160. var resp region_sync.BatchDeductResponse
  161. json.Unmarshal(w.Body.Bytes(), &resp)
  162. assert.True(t, resp.Success)
  163. assert.Len(t, resp.Results, 2)
  164. // 验证用户余额已更新
  165. var updatedUser1 model.User
  166. db.First(&updatedUser1, user1.Id)
  167. assert.Equal(t, 900, updatedUser1.Quota)
  168. var updatedUser2 model.User
  169. db.First(&updatedUser2, user2.Id)
  170. assert.Equal(t, 450, updatedUser2.Quota)
  171. }
  172. func TestBatchDeductQuota_InsufficientQuota(t *testing.T) {
  173. gin.SetMode(gin.TestMode)
  174. router := gin.New()
  175. router.Use(testSyncAuthMiddleware())
  176. router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota)
  177. origDB := model.DB
  178. defer func() {
  179. model.DB = origDB
  180. }()
  181. // 设置测试数据库
  182. db := setupSyncTestDB(t)
  183. model.DB = db
  184. // 创建余额不足的用户
  185. user := model.User{Username: "pooruser", Quota: 10, AffCode: "aff-poor-1"}
  186. db.Create(&user)
  187. req := region_sync.BatchDeductRequest{
  188. Records: []region_sync.BatchDeductRecord{
  189. {UserId: user.Id, RequestId: "req-1", Quota: 100},
  190. },
  191. }
  192. reqBody, _ := json.Marshal(req)
  193. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody))
  194. reqHTTP.Header.Set("Content-Type", "application/json")
  195. w := httptest.NewRecorder()
  196. router.ServeHTTP(w, reqHTTP)
  197. assert.Equal(t, http.StatusOK, w.Code)
  198. var resp region_sync.BatchDeductResponse
  199. json.Unmarshal(w.Body.Bytes(), &resp)
  200. assert.True(t, resp.Success)
  201. assert.Len(t, resp.Results, 1)
  202. assert.False(t, resp.Results[0].Success)
  203. assert.Contains(t, resp.Results[0].Error, "insufficient quota")
  204. }
  205. // testSyncAuthMiddleware 测试专用的认证中间件,绕过系统设置
  206. func testSyncAuthMiddleware() gin.HandlerFunc {
  207. return func(c *gin.Context) {
  208. c.Set("sync_node", "test-node")
  209. c.Next()
  210. }
  211. }
  212. // setupSyncTestDB 设置同步测试数据库
  213. func setupSyncTestDB(t *testing.T) *gorm.DB {
  214. t.Helper()
  215. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  216. require.NoError(t, err)
  217. sqlDB, _ := db.DB()
  218. sqlDB.SetMaxOpenConns(1)
  219. err = db.AutoMigrate(&model.User{}, &model.QuotaSyncLog{})
  220. require.NoError(t, err)
  221. return db
  222. }
  223. // ---------------------------------------------------------------------------
  224. // 阶段 4:Controller 边界条件测试
  225. // ---------------------------------------------------------------------------
  226. func TestReceiveSyncedUserCreate_DuplicateUser(t *testing.T) {
  227. gin.SetMode(gin.TestMode)
  228. router := gin.New()
  229. router.Use(testSyncAuthMiddleware())
  230. router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate)
  231. origDB := model.DB
  232. defer func() { model.DB = origDB }()
  233. db := setupSyncTestDB(t)
  234. model.DB = db
  235. // 先创建一个用户
  236. db.Create(&model.User{
  237. Username: "dupuser", RemoteUserId: 10000001,
  238. Source: common.UserSourceSynced, Password: "hash", AffCode: "aff-dup-1",
  239. })
  240. req := region_sync.SyncUserRequest{
  241. Username: "dupuser",
  242. RemoteUserId: 10000001,
  243. PasswordHash: "newhash",
  244. }
  245. reqBody, _ := json.Marshal(req)
  246. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", bytes.NewBuffer(reqBody))
  247. reqHTTP.Header.Set("Content-Type", "application/json")
  248. w := httptest.NewRecorder()
  249. router.ServeHTTP(w, reqHTTP)
  250. assert.Equal(t, http.StatusOK, w.Code)
  251. var resp map[string]interface{}
  252. json.Unmarshal(w.Body.Bytes(), &resp)
  253. assert.True(t, resp["success"].(bool))
  254. assert.Contains(t, resp["message"], "already exists")
  255. }
  256. func TestReceiveSyncedUserCreate_InvalidBody(t *testing.T) {
  257. gin.SetMode(gin.TestMode)
  258. router := gin.New()
  259. router.Use(testSyncAuthMiddleware())
  260. router.POST("/api/internal/sync/user/create", ReceiveSyncedUserCreate)
  261. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/user/create", strings.NewReader("not json"))
  262. reqHTTP.Header.Set("Content-Type", "application/json")
  263. w := httptest.NewRecorder()
  264. router.ServeHTTP(w, reqHTTP)
  265. assert.Equal(t, http.StatusBadRequest, w.Code)
  266. }
  267. func TestReceiveQuotaUpdate_UserNotFound(t *testing.T) {
  268. gin.SetMode(gin.TestMode)
  269. router := gin.New()
  270. router.Use(testSyncAuthMiddleware())
  271. router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate)
  272. origDB := model.DB
  273. defer func() { model.DB = origDB }()
  274. db := setupSyncTestDB(t)
  275. model.DB = db
  276. req := region_sync.UpdateQuotaRequest{
  277. RemoteUserId: 99999999,
  278. Quota: 5000,
  279. }
  280. reqBody, _ := json.Marshal(req)
  281. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", bytes.NewBuffer(reqBody))
  282. reqHTTP.Header.Set("Content-Type", "application/json")
  283. w := httptest.NewRecorder()
  284. router.ServeHTTP(w, reqHTTP)
  285. assert.Equal(t, http.StatusNotFound, w.Code)
  286. }
  287. func TestReceiveQuotaUpdate_InvalidBody(t *testing.T) {
  288. gin.SetMode(gin.TestMode)
  289. router := gin.New()
  290. router.Use(testSyncAuthMiddleware())
  291. router.POST("/api/internal/sync/quota/update", ReceiveQuotaUpdate)
  292. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/update", strings.NewReader("invalid"))
  293. reqHTTP.Header.Set("Content-Type", "application/json")
  294. w := httptest.NewRecorder()
  295. router.ServeHTTP(w, reqHTTP)
  296. assert.Equal(t, http.StatusBadRequest, w.Code)
  297. }
  298. func TestQueryUserQuota_UserNotFound(t *testing.T) {
  299. gin.SetMode(gin.TestMode)
  300. router := gin.New()
  301. router.Use(testSyncAuthMiddleware())
  302. router.POST("/api/internal/sync/quota/query", QueryUserQuota)
  303. origDB := model.DB
  304. defer func() { model.DB = origDB }()
  305. db := setupSyncTestDB(t)
  306. model.DB = db
  307. req := region_sync.QueryQuotaRequest{UserId: 99999}
  308. reqBody, _ := json.Marshal(req)
  309. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/query", bytes.NewBuffer(reqBody))
  310. reqHTTP.Header.Set("Content-Type", "application/json")
  311. w := httptest.NewRecorder()
  312. router.ServeHTTP(w, reqHTTP)
  313. assert.Equal(t, http.StatusNotFound, w.Code)
  314. }
  315. func TestBatchDeductQuota_EmptyRecords(t *testing.T) {
  316. gin.SetMode(gin.TestMode)
  317. router := gin.New()
  318. router.Use(testSyncAuthMiddleware())
  319. router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota)
  320. origDB := model.DB
  321. defer func() { model.DB = origDB }()
  322. db := setupSyncTestDB(t)
  323. model.DB = db
  324. req := region_sync.BatchDeductRequest{Records: []region_sync.BatchDeductRecord{}}
  325. reqBody, _ := json.Marshal(req)
  326. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody))
  327. reqHTTP.Header.Set("Content-Type", "application/json")
  328. w := httptest.NewRecorder()
  329. router.ServeHTTP(w, reqHTTP)
  330. assert.Equal(t, http.StatusOK, w.Code)
  331. var resp region_sync.BatchDeductResponse
  332. json.Unmarshal(w.Body.Bytes(), &resp)
  333. assert.True(t, resp.Success)
  334. assert.Len(t, resp.Results, 0)
  335. }
  336. func TestBatchDeductQuota_UserNotFound(t *testing.T) {
  337. gin.SetMode(gin.TestMode)
  338. router := gin.New()
  339. router.Use(testSyncAuthMiddleware())
  340. router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota)
  341. origDB := model.DB
  342. defer func() { model.DB = origDB }()
  343. db := setupSyncTestDB(t)
  344. model.DB = db
  345. req := region_sync.BatchDeductRequest{
  346. Records: []region_sync.BatchDeductRecord{
  347. {UserId: 99999, RequestId: "req-nf", Quota: 100},
  348. },
  349. }
  350. reqBody, _ := json.Marshal(req)
  351. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", bytes.NewBuffer(reqBody))
  352. reqHTTP.Header.Set("Content-Type", "application/json")
  353. w := httptest.NewRecorder()
  354. router.ServeHTTP(w, reqHTTP)
  355. assert.Equal(t, http.StatusOK, w.Code)
  356. var resp region_sync.BatchDeductResponse
  357. json.Unmarshal(w.Body.Bytes(), &resp)
  358. assert.True(t, resp.Success)
  359. assert.Len(t, resp.Results, 1)
  360. assert.False(t, resp.Results[0].Success)
  361. assert.Contains(t, resp.Results[0].Error, "user not found")
  362. }
  363. func TestBatchDeductQuota_InvalidBody(t *testing.T) {
  364. gin.SetMode(gin.TestMode)
  365. router := gin.New()
  366. router.Use(testSyncAuthMiddleware())
  367. router.POST("/api/internal/sync/quota/batch-deduct", BatchDeductQuota)
  368. reqHTTP, _ := http.NewRequest("POST", "/api/internal/sync/quota/batch-deduct", strings.NewReader("bad"))
  369. reqHTTP.Header.Set("Content-Type", "application/json")
  370. w := httptest.NewRecorder()
  371. router.ServeHTTP(w, reqHTTP)
  372. assert.Equal(t, http.StatusBadRequest, w.Code)
  373. }