|
- package region_sync
-
- import (
- "encoding/json"
- "io"
- "net/http"
- "net/http/httptest"
- "sync/atomic"
- "testing"
- "time"
-
- "github.com/QuantumNous/new-api/model"
- "github.com/QuantumNous/new-api/setting/system_setting"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupPushDB(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 := model.DB
- origLogDB := model.LOG_DB
- model.DB = db
- model.LOG_DB = db
-
- require.NoError(t, db.AutoMigrate(&model.QuotaSyncLog{}))
-
- t.Cleanup(func() {
- model.DB = origDB
- model.LOG_DB = origLogDB
- sqlDB.Close()
- })
- return db
- }
-
- func setupPushSettings(t *testing.T, fn func(s *system_setting.RegionSyncSettings)) {
- t.Helper()
- s := system_setting.GetRegionSyncSettings()
- orig := *s
- fn(s)
- t.Cleanup(func() { *s = orig })
- }
-
- func TestPushUserCreateToSlave_Disabled(t *testing.T) {
- setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
- s.Enabled = false
- })
-
- // 不应该 panic 或发请求
- PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
- }
-
- func TestPushUserCreateToSlave_NotMaster(t *testing.T) {
- setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
- s.Enabled = true
- s.IsMaster = false
- })
-
- PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
- }
-
- func TestPushUserCreateToSlave_Success(t *testing.T) {
- setupPushDB(t)
-
- var receivedBody []byte
- var receivedKey string
- server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- receivedKey = r.Header.Get("X-Sync-API-Key")
- receivedBody, _ = io.ReadAll(r.Body)
- w.Header().Set("Content-Type", "application/json")
- json.NewEncoder(w).Encode(SyncUserResponse{Success: true, Message: "ok"})
- }))
- defer server.Close()
-
- setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
- s.Enabled = true
- s.IsMaster = true
- s.SlaveEndpoints = []string{server.URL}
- s.SyncApiKey = "test-key-123"
- })
-
- PushUserCreateToSlave(&model.User{
- Id: 100, Username: "syncuser", Email: "a@b.com",
- Password: "hashed", DisplayName: "Sync User", Quota: 500000, Group: "default",
- })
-
- // 等待异步 goroutine 完成
- time.Sleep(200 * time.Millisecond)
-
- assert.Equal(t, "test-key-123", receivedKey)
- var req SyncUserRequest
- require.NoError(t, json.Unmarshal(receivedBody, &req))
- assert.Equal(t, "syncuser", req.Username)
- assert.Equal(t, "hashed", req.PasswordHash)
- assert.Equal(t, 100, req.RemoteUserId)
- assert.Equal(t, 500000, req.Quota)
- }
-
- func TestPushUserCreateToSlave_PartialFailure(t *testing.T) {
- setupPushDB(t)
-
- var successCount int32
- server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- atomic.AddInt32(&successCount, 1)
- w.Header().Set("Content-Type", "application/json")
- json.NewEncoder(w).Encode(SyncUserResponse{Success: true})
- }))
- defer server1.Close()
-
- server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- w.Header().Set("Content-Type", "application/json")
- json.NewEncoder(w).Encode(SyncUserResponse{Success: false, Error: "rejected"})
- }))
- defer server2.Close()
-
- setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
- s.Enabled = true
- s.IsMaster = true
- s.SlaveEndpoints = []string{server1.URL, server2.URL}
- s.SyncApiKey = "test-key"
- })
-
- PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
- time.Sleep(200 * time.Millisecond)
-
- assert.Equal(t, int32(1), atomic.LoadInt32(&successCount))
- }
-
- func TestPushQuotaUpdateToSlave_Success(t *testing.T) {
- var receivedBody []byte
- server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- receivedBody, _ = io.ReadAll(r.Body)
- w.Header().Set("Content-Type", "application/json")
- json.NewEncoder(w).Encode(UpdateQuotaResponse{Success: true})
- }))
- defer server.Close()
-
- setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
- s.Enabled = true
- s.IsMaster = true
- s.SlaveEndpoints = []string{server.URL}
- s.SyncApiKey = "key"
- })
-
- PushQuotaUpdateToSlave(100, 500000)
- time.Sleep(200 * time.Millisecond)
-
- var req UpdateQuotaRequest
- require.NoError(t, json.Unmarshal(receivedBody, &req))
- assert.Equal(t, 100, req.RemoteUserId)
- assert.Equal(t, 500000, req.Quota)
- }
-
- func TestPushUserCreateToSlave_NoEndpoints(t *testing.T) {
- setupPushDB(t)
- setupPushSettings(t, func(s *system_setting.RegionSyncSettings) {
- s.Enabled = true
- s.IsMaster = true
- s.SlaveEndpoints = nil
- s.SyncApiKey = "key"
- })
-
- // 不应该 panic
- PushUserCreateToSlave(&model.User{Id: 1, Username: "test"})
- }
|