Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

492 строки
14 KiB

  1. package controller
  2. import (
  3. "encoding/json"
  4. "net/http"
  5. "net/http/httptest"
  6. "strings"
  7. "testing"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/model"
  10. user_migration "github.com/QuantumNous/new-api/service/user_migration"
  11. "github.com/gin-gonic/gin"
  12. "github.com/glebarez/sqlite"
  13. "github.com/stretchr/testify/require"
  14. "gorm.io/gorm"
  15. )
  16. func setupInternalMigrationDB(t *testing.T) *gorm.DB {
  17. t.Helper()
  18. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  19. require.NoError(t, err)
  20. sqlDB, err := db.DB()
  21. require.NoError(t, err)
  22. sqlDB.SetMaxOpenConns(1)
  23. require.NoError(t, db.AutoMigrate(&model.User{}, &model.CustomOAuthProvider{}, &model.UserOAuthBinding{}))
  24. return db
  25. }
  26. func TestListMigrationUsers_ReturnsLocalOnly(t *testing.T) {
  27. gin.SetMode(gin.TestMode)
  28. router := gin.New()
  29. router.Use(testSyncAuthMiddleware())
  30. router.GET("/api/internal/migration/users", ListMigrationUsers)
  31. db := setupInternalMigrationDB(t)
  32. orig := model.DB
  33. model.DB = db
  34. defer func() {
  35. model.DB = orig
  36. }()
  37. require.NoError(t, db.Create(&model.User{
  38. Id: 10000001,
  39. Username: "ov-1",
  40. Password: "hash",
  41. Email: "u1@example.com",
  42. AffCode: "A1",
  43. Source: common.UserSourceLocal,
  44. Quota: 99,
  45. }).Error)
  46. require.NoError(t, db.Create(&model.User{
  47. Id: 10000002,
  48. Username: "ov-2",
  49. Password: "hash",
  50. AffCode: "A2",
  51. Source: common.UserSourceSynced,
  52. RemoteUserId: 1,
  53. }).Error)
  54. require.NoError(t, db.Create(&model.User{
  55. Id: 10000003,
  56. Username: "root",
  57. Password: "hash",
  58. AffCode: "A3",
  59. Source: common.UserSourceLocal,
  60. Role: common.RoleRootUser,
  61. }).Error)
  62. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users?page=0&page_size=20", nil)
  63. w := httptest.NewRecorder()
  64. router.ServeHTTP(w, req)
  65. require.Equal(t, http.StatusOK, w.Code)
  66. var body map[string]any
  67. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  68. require.True(t, body["success"].(bool))
  69. require.EqualValues(t, 1, body["total"])
  70. }
  71. func TestListMigrationUsers_FiltersByKeyword(t *testing.T) {
  72. gin.SetMode(gin.TestMode)
  73. router := gin.New()
  74. router.Use(testSyncAuthMiddleware())
  75. router.GET("/api/internal/migration/users", ListMigrationUsers)
  76. db := setupInternalMigrationDB(t)
  77. orig := model.DB
  78. model.DB = db
  79. defer func() {
  80. model.DB = orig
  81. }()
  82. require.NoError(t, db.Create(&model.User{
  83. Id: 10000011,
  84. Username: "alice-local",
  85. Password: "hash",
  86. Email: "alice@example.com",
  87. AffCode: "A11",
  88. Source: common.UserSourceLocal,
  89. }).Error)
  90. require.NoError(t, db.Create(&model.User{
  91. Id: 10000012,
  92. Username: "bob-local",
  93. Password: "hash",
  94. Email: "bob@example.com",
  95. AffCode: "A12",
  96. Source: common.UserSourceLocal,
  97. }).Error)
  98. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users?page=0&page_size=20&keyword=alice", nil)
  99. w := httptest.NewRecorder()
  100. router.ServeHTTP(w, req)
  101. require.Equal(t, http.StatusOK, w.Code)
  102. var body struct {
  103. Success bool `json:"success"`
  104. Data []user_migration.RemoteUserSnapshot `json:"data"`
  105. Total int `json:"total"`
  106. }
  107. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  108. require.True(t, body.Success)
  109. require.Equal(t, 1, body.Total)
  110. require.Len(t, body.Data, 1)
  111. require.Equal(t, 10000011, body.Data[0].Id)
  112. }
  113. func TestQueryMigrationUsers_OnlyReturnsSelectedLocalUsers(t *testing.T) {
  114. gin.SetMode(gin.TestMode)
  115. router := gin.New()
  116. router.Use(testSyncAuthMiddleware())
  117. router.POST("/api/internal/migration/users/query", QueryMigrationUsers)
  118. db := setupInternalMigrationDB(t)
  119. orig := model.DB
  120. model.DB = db
  121. defer func() {
  122. model.DB = orig
  123. }()
  124. require.NoError(t, db.Create(&model.User{Id: 10000001, Username: "u1", Password: "hash", AffCode: "A1", Source: common.UserSourceLocal}).Error)
  125. require.NoError(t, db.Create(&model.User{Id: 10000002, Username: "u2", Password: "hash", AffCode: "A2", Source: common.UserSourceLocal, Role: common.RoleRootUser}).Error)
  126. require.NoError(t, db.Create(&model.User{Id: 10000003, Username: "u3", Password: "hash", AffCode: "A3", Source: common.UserSourceSynced}).Error)
  127. req := httptest.NewRequest(http.MethodPost, "/api/internal/migration/users/query", strings.NewReader(`{"page":0,"page_size":50,"selection_mode":"explicit_ids","source_user_ids":[10000001,10000002,10000003,10000004]}`))
  128. req.Header.Set("Content-Type", "application/json")
  129. w := httptest.NewRecorder()
  130. router.ServeHTTP(w, req)
  131. require.Equal(t, http.StatusOK, w.Code)
  132. var body struct {
  133. Success bool `json:"success"`
  134. Data []user_migration.RemoteUserSnapshot `json:"data"`
  135. Total int `json:"total"`
  136. SelectionSummary user_migration.SelectionSummary `json:"selection_summary"`
  137. }
  138. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  139. require.True(t, body.Success)
  140. require.Len(t, body.Data, 1)
  141. require.Equal(t, 10000001, body.Data[0].Id)
  142. require.Equal(t, 1, body.Total)
  143. require.Equal(t, 4, body.SelectionSummary.Requested)
  144. require.Equal(t, 1, body.SelectionSummary.Matched)
  145. require.Len(t, body.SelectionSummary.Excluded, 3)
  146. }
  147. func TestQueryMigrationUsers_InvalidSelectionModeRejected(t *testing.T) {
  148. gin.SetMode(gin.TestMode)
  149. router := gin.New()
  150. router.Use(testSyncAuthMiddleware())
  151. router.POST("/api/internal/migration/users/query", QueryMigrationUsers)
  152. req := httptest.NewRequest(http.MethodPost, "/api/internal/migration/users/query", strings.NewReader(`{"page":0,"page_size":50,"selection_mode":"all","source_user_ids":[1]}`))
  153. req.Header.Set("Content-Type", "application/json")
  154. w := httptest.NewRecorder()
  155. router.ServeHTTP(w, req)
  156. require.Equal(t, http.StatusBadRequest, w.Code)
  157. }
  158. func TestGetMigrationUser_RootExcluded(t *testing.T) {
  159. gin.SetMode(gin.TestMode)
  160. router := gin.New()
  161. router.Use(testSyncAuthMiddleware())
  162. router.GET("/api/internal/migration/users/:id", GetMigrationUser)
  163. db := setupInternalMigrationDB(t)
  164. orig := model.DB
  165. model.DB = db
  166. defer func() {
  167. model.DB = orig
  168. }()
  169. require.NoError(t, db.Create(&model.User{
  170. Id: 10000011,
  171. Username: "root",
  172. Password: "hash",
  173. AffCode: "ROOT1",
  174. Source: common.UserSourceLocal,
  175. Role: common.RoleRootUser,
  176. }).Error)
  177. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/10000011", nil)
  178. w := httptest.NewRecorder()
  179. router.ServeHTTP(w, req)
  180. require.Equal(t, http.StatusNotFound, w.Code)
  181. }
  182. func TestGetMigrationUser_SyncedUserStillReadableForVerify(t *testing.T) {
  183. gin.SetMode(gin.TestMode)
  184. router := gin.New()
  185. router.Use(testSyncAuthMiddleware())
  186. router.GET("/api/internal/migration/users/:id", GetMigrationUser)
  187. db := setupInternalMigrationDB(t)
  188. orig := model.DB
  189. model.DB = db
  190. defer func() {
  191. model.DB = orig
  192. }()
  193. require.NoError(t, db.Create(&model.User{
  194. Id: 10000021,
  195. Username: "synced-user",
  196. Password: "hash",
  197. AffCode: "SYNC1",
  198. Source: common.UserSourceSynced,
  199. RemoteUserId: 88,
  200. SyncedQuota: 123,
  201. Role: common.RoleCommonUser,
  202. }).Error)
  203. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/10000021", nil)
  204. w := httptest.NewRecorder()
  205. router.ServeHTTP(w, req)
  206. require.Equal(t, http.StatusOK, w.Code)
  207. var body map[string]any
  208. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  209. require.True(t, body["success"].(bool))
  210. }
  211. func TestGetMigrationUser_InvalidID(t *testing.T) {
  212. gin.SetMode(gin.TestMode)
  213. router := gin.New()
  214. router.Use(testSyncAuthMiddleware())
  215. router.GET("/api/internal/migration/users/:id", GetMigrationUser)
  216. req := httptest.NewRequest(http.MethodGet, "/api/internal/migration/users/not-a-number", nil)
  217. w := httptest.NewRecorder()
  218. router.ServeHTTP(w, req)
  219. require.Equal(t, http.StatusBadRequest, w.Code)
  220. }
  221. func TestConvertMigrationUserToSynced_Idempotent(t *testing.T) {
  222. gin.SetMode(gin.TestMode)
  223. router := gin.New()
  224. router.Use(testSyncAuthMiddleware())
  225. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  226. db := setupInternalMigrationDB(t)
  227. orig := model.DB
  228. model.DB = db
  229. defer func() {
  230. model.DB = orig
  231. }()
  232. require.NoError(t, db.Create(&model.User{
  233. Id: 10000002,
  234. Username: "ov-2",
  235. Password: "hash",
  236. AffCode: "A2",
  237. Source: common.UserSourceLocal,
  238. Quota: 66,
  239. }).Error)
  240. body := `{"remote_user_id":2002,"synced_quota":66}`
  241. for i := 0; i < 2; i++ {
  242. req := httptest.NewRequest(http.MethodPost,
  243. "/api/internal/migration/users/10000002/convert-to-synced",
  244. strings.NewReader(body))
  245. req.Header.Set("Content-Type", "application/json")
  246. w := httptest.NewRecorder()
  247. router.ServeHTTP(w, req)
  248. require.Equal(t, http.StatusOK, w.Code)
  249. }
  250. }
  251. func TestConvertMigrationUserToSynced_ConflictRemoteId(t *testing.T) {
  252. gin.SetMode(gin.TestMode)
  253. router := gin.New()
  254. router.Use(testSyncAuthMiddleware())
  255. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  256. db := setupInternalMigrationDB(t)
  257. orig := model.DB
  258. model.DB = db
  259. defer func() {
  260. model.DB = orig
  261. }()
  262. require.NoError(t, db.Create(&model.User{
  263. Id: 10000003,
  264. Username: "ov-3",
  265. Password: "hash",
  266. AffCode: "A3",
  267. Source: common.UserSourceSynced,
  268. RemoteUserId: 5000,
  269. }).Error)
  270. body := `{"remote_user_id":9999,"synced_quota":100}`
  271. req := httptest.NewRequest(http.MethodPost,
  272. "/api/internal/migration/users/10000003/convert-to-synced",
  273. strings.NewReader(body))
  274. req.Header.Set("Content-Type", "application/json")
  275. w := httptest.NewRecorder()
  276. router.ServeHTTP(w, req)
  277. require.Equal(t, http.StatusConflict, w.Code)
  278. }
  279. func TestConvertMigrationUserToSynced_ConflictUsesSentinelErrors(t *testing.T) {
  280. gin.SetMode(gin.TestMode)
  281. router := gin.New()
  282. router.Use(testSyncAuthMiddleware())
  283. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  284. db := setupInternalMigrationDB(t)
  285. orig := model.DB
  286. model.DB = db
  287. defer func() {
  288. model.DB = orig
  289. }()
  290. require.NoError(t, db.Create(&model.User{
  291. Id: 10000013,
  292. Username: "ov-13",
  293. Password: "hash",
  294. AffCode: "A13",
  295. Source: common.UserSourceSynced,
  296. RemoteUserId: 5001,
  297. }).Error)
  298. body := `{"remote_user_id":9999,"synced_quota":100}`
  299. req := httptest.NewRequest(http.MethodPost,
  300. "/api/internal/migration/users/10000013/convert-to-synced",
  301. strings.NewReader(body))
  302. req.Header.Set("Content-Type", "application/json")
  303. w := httptest.NewRecorder()
  304. router.ServeHTTP(w, req)
  305. require.Equal(t, http.StatusConflict, w.Code)
  306. }
  307. func TestConvertMigrationUserToSynced_RejectsRootUser(t *testing.T) {
  308. gin.SetMode(gin.TestMode)
  309. router := gin.New()
  310. router.Use(testSyncAuthMiddleware())
  311. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  312. db := setupInternalMigrationDB(t)
  313. orig := model.DB
  314. model.DB = db
  315. defer func() {
  316. model.DB = orig
  317. }()
  318. require.NoError(t, db.Create(&model.User{
  319. Id: 10000012,
  320. Username: "root",
  321. Password: "hash",
  322. AffCode: "ROOT2",
  323. Source: common.UserSourceLocal,
  324. Role: common.RoleRootUser,
  325. }).Error)
  326. body := `{"remote_user_id":9999,"synced_quota":100}`
  327. req := httptest.NewRequest(http.MethodPost,
  328. "/api/internal/migration/users/10000012/convert-to-synced",
  329. strings.NewReader(body))
  330. req.Header.Set("Content-Type", "application/json")
  331. w := httptest.NewRecorder()
  332. router.ServeHTTP(w, req)
  333. require.Equal(t, http.StatusConflict, w.Code)
  334. }
  335. func TestConvertMigrationUserToSynced_InvalidRemoteUserID(t *testing.T) {
  336. gin.SetMode(gin.TestMode)
  337. router := gin.New()
  338. router.Use(testSyncAuthMiddleware())
  339. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  340. db := setupInternalMigrationDB(t)
  341. orig := model.DB
  342. model.DB = db
  343. defer func() {
  344. model.DB = orig
  345. }()
  346. require.NoError(t, db.Create(&model.User{
  347. Id: 10000014,
  348. Username: "ov-14",
  349. Password: "hash",
  350. AffCode: "A14",
  351. Source: common.UserSourceLocal,
  352. }).Error)
  353. body := `{"remote_user_id":0,"synced_quota":100}`
  354. req := httptest.NewRequest(http.MethodPost,
  355. "/api/internal/migration/users/10000014/convert-to-synced",
  356. strings.NewReader(body))
  357. req.Header.Set("Content-Type", "application/json")
  358. w := httptest.NewRecorder()
  359. router.ServeHTTP(w, req)
  360. require.Equal(t, http.StatusBadRequest, w.Code)
  361. }
  362. func TestConvertMigrationUserToSynced_PersistsSyncedFields(t *testing.T) {
  363. gin.SetMode(gin.TestMode)
  364. router := gin.New()
  365. router.Use(testSyncAuthMiddleware())
  366. router.POST("/api/internal/migration/users/:id/convert-to-synced", ConvertMigrationUserToSynced)
  367. db := setupInternalMigrationDB(t)
  368. orig := model.DB
  369. model.DB = db
  370. defer func() {
  371. model.DB = orig
  372. }()
  373. require.NoError(t, db.Create(&model.User{
  374. Id: 10000015,
  375. Username: "ov-15",
  376. Password: "hash",
  377. AffCode: "A15",
  378. Source: common.UserSourceLocal,
  379. Quota: 88,
  380. }).Error)
  381. body := `{"remote_user_id":2015,"synced_quota":166}`
  382. req := httptest.NewRequest(http.MethodPost,
  383. "/api/internal/migration/users/10000015/convert-to-synced",
  384. strings.NewReader(body))
  385. req.Header.Set("Content-Type", "application/json")
  386. w := httptest.NewRecorder()
  387. router.ServeHTTP(w, req)
  388. require.Equal(t, http.StatusOK, w.Code)
  389. var saved model.User
  390. require.NoError(t, db.First(&saved, 10000015).Error)
  391. require.Equal(t, common.UserSourceSynced, saved.Source)
  392. require.Equal(t, 2015, saved.RemoteUserId)
  393. require.Equal(t, 166, saved.SyncedQuota)
  394. require.Greater(t, saved.LastSyncAt, int64(0))
  395. }
  396. func TestCheckSyncedCopy(t *testing.T) {
  397. gin.SetMode(gin.TestMode)
  398. router := gin.New()
  399. router.Use(testSyncAuthMiddleware())
  400. router.GET("/api/internal/migration/users/check-synced-copy", CheckSyncedCopy)
  401. db := setupInternalMigrationDB(t)
  402. orig := model.DB
  403. model.DB = db
  404. defer func() {
  405. model.DB = orig
  406. }()
  407. require.NoError(t, db.Create(&model.User{
  408. Id: 10000010,
  409. Username: "synced",
  410. Password: "hash",
  411. AffCode: "S1",
  412. Source: common.UserSourceSynced,
  413. RemoteUserId: 456,
  414. }).Error)
  415. req := httptest.NewRequest(http.MethodGet,
  416. "/api/internal/migration/users/check-synced-copy?cn_user_id=456", nil)
  417. w := httptest.NewRecorder()
  418. router.ServeHTTP(w, req)
  419. require.Equal(t, http.StatusOK, w.Code)
  420. var body map[string]any
  421. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  422. require.True(t, body["has_copy"].(bool))
  423. }