Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

172 Zeilen
4.6 KiB

  1. package region_sync
  2. import (
  3. "encoding/json"
  4. "io"
  5. "net/http"
  6. "net/http/httptest"
  7. "sync/atomic"
  8. "testing"
  9. "time"
  10. "github.com/QuantumNous/new-api/model"
  11. "github.com/QuantumNous/new-api/setting/system_setting"
  12. "github.com/glebarez/sqlite"
  13. "github.com/stretchr/testify/assert"
  14. "github.com/stretchr/testify/require"
  15. "gorm.io/gorm"
  16. )
  17. func setupPushDB(t *testing.T) *gorm.DB {
  18. t.Helper()
  19. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  20. require.NoError(t, err)
  21. sqlDB, _ := db.DB()
  22. sqlDB.SetMaxOpenConns(1)
  23. origDB := model.DB
  24. origLogDB := model.LOG_DB
  25. model.DB = db
  26. model.LOG_DB = db
  27. require.NoError(t, db.AutoMigrate(&model.QuotaSyncLog{}))
  28. t.Cleanup(func() {
  29. model.DB = origDB
  30. model.LOG_DB = origLogDB
  31. sqlDB.Close()
  32. })
  33. return db
  34. }
  35. func setupPushSettings(t *testing.T, fn func(s *system_setting.RegionSyncSettings)) {
  36. t.Helper()
  37. s := system_setting.GetRegionSyncSettings()
  38. orig := *s
  39. fn(s)
  40. t.Cleanup(func() { *s = orig })
  41. }
  42. func TestPushUserCreateToSlave_Disabled(t *testing.T) {
  43. setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
  44. s.Enabled = false
  45. })
  46. // 不应该 panic 或发请求
  47. PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
  48. }
  49. func TestPushUserCreateToSlave_NotMaster(t *testing.T) {
  50. setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
  51. s.Enabled = true
  52. s.IsMaster = false
  53. })
  54. PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
  55. }
  56. func TestPushUserCreateToSlave_Success(t *testing.T) {
  57. setupPushDB(t)
  58. var receivedBody []byte
  59. var receivedKey string
  60. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  61. receivedKey = r.Header.Get("X-Sync-API-Key")
  62. receivedBody, _ = io.ReadAll(r.Body)
  63. w.Header().Set("Content-Type", "application/json")
  64. json.NewEncoder(w).Encode(SyncUserResponse{Success: true, Message: "ok"})
  65. }))
  66. defer server.Close()
  67. setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
  68. s.Enabled = true
  69. s.IsMaster = true
  70. s.SlaveEndpoints = []string{server.URL}
  71. s.SyncApiKey = "test-key-123"
  72. })
  73. PushUserCreateToSlave(&model.User{
  74. Id: 100, Username: "syncuser", Email: "a@b.com",
  75. Password: "hashed", DisplayName: "Sync User", Quota: 500000, Group: "default",
  76. })
  77. // 等待异步 goroutine 完成
  78. time.Sleep(200 * time.Millisecond)
  79. assert.Equal(t, "test-key-123", receivedKey)
  80. var req SyncUserRequest
  81. require.NoError(t, json.Unmarshal(receivedBody, &req))
  82. assert.Equal(t, "syncuser", req.Username)
  83. assert.Equal(t, "hashed", req.PasswordHash)
  84. assert.Equal(t, 100, req.RemoteUserId)
  85. assert.Equal(t, 500000, req.Quota)
  86. }
  87. func TestPushUserCreateToSlave_PartialFailure(t *testing.T) {
  88. setupPushDB(t)
  89. var successCount int32
  90. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  91. atomic.AddInt32(&successCount, 1)
  92. w.Header().Set("Content-Type", "application/json")
  93. json.NewEncoder(w).Encode(SyncUserResponse{Success: true})
  94. }))
  95. defer server1.Close()
  96. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  97. w.Header().Set("Content-Type", "application/json")
  98. json.NewEncoder(w).Encode(SyncUserResponse{Success: false, Error: "rejected"})
  99. }))
  100. defer server2.Close()
  101. setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
  102. s.Enabled = true
  103. s.IsMaster = true
  104. s.SlaveEndpoints = []string{server1.URL, server2.URL}
  105. s.SyncApiKey = "test-key"
  106. })
  107. PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
  108. time.Sleep(200 * time.Millisecond)
  109. assert.Equal(t, int32(1), atomic.LoadInt32(&successCount))
  110. }
  111. func TestPushQuotaUpdateToSlave_Success(t *testing.T) {
  112. var receivedBody []byte
  113. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  114. receivedBody, _ = io.ReadAll(r.Body)
  115. w.Header().Set("Content-Type", "application/json")
  116. json.NewEncoder(w).Encode(UpdateQuotaResponse{Success: true})
  117. }))
  118. defer server.Close()
  119. setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
  120. s.Enabled = true
  121. s.IsMaster = true
  122. s.SlaveEndpoints = []string{server.URL}
  123. s.SyncApiKey = "key"
  124. })
  125. PushQuotaUpdateToSlave(100, 500000)
  126. time.Sleep(200 * time.Millisecond)
  127. var req UpdateQuotaRequest
  128. require.NoError(t, json.Unmarshal(receivedBody, &req))
  129. assert.Equal(t, 100, req.RemoteUserId)
  130. assert.Equal(t, 500000, req.Quota)
  131. }
  132. func TestPushUserCreateToSlave_NoEndpoints(t *testing.T) {
  133. setupPushDB(t)
  134. setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
  135. s.Enabled = true
  136. s.IsMaster = true
  137. s.SlaveEndpoints = nil
  138. s.SyncApiKey = "key"
  139. })
  140. // 不应该 panic
  141. PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
  142. }