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"}) }