package helpers import ( "fmt" "net/http" "time" "github.com/QuantumNous/new-api/common" ) type optionResponse struct { Success bool `json:"success"` Message string `json:"message,omitempty"` } type internalListUsersResponse struct { Success bool `json:"success"` Total int `json:"total"` } func SetOption(client *APIClient, key string, value any) error { var resp optionResponse if err := client.PutJSON("/api/option/", map[string]any{ "key": key, "value": value, }, http.StatusOK, &resp); err != nil { return err } if !resp.Success { return fmt.Errorf("set option %s failed: %s", key, resp.Message) } return nil } func ConfigureCNRegionSync(client *APIClient, cnBaseURL, ovBaseURL, apiKey string) error { slaveEndpoints, err := common.Marshal([]string{ovBaseURL}) if err != nil { return fmt.Errorf("marshal slave endpoints: %w", err) } options := []struct { Key string Value any }{ {Key: "region_sync.enabled", Value: true}, {Key: "region_sync.region_id", Value: "cn"}, {Key: "region_sync.is_master", Value: true}, {Key: "region_sync.master_endpoint", Value: cnBaseURL}, {Key: "region_sync.slave_endpoints", Value: string(slaveEndpoints)}, {Key: "region_sync.sync_api_key", Value: apiKey}, } for _, option := range options { if err := SetOption(client, option.Key, option.Value); err != nil { return err } } return nil } func ConfigureOVRegionSync(client *APIClient, cnBaseURL, apiKey string) error { emptyEndpoints, err := common.Marshal([]string{}) if err != nil { return fmt.Errorf("marshal empty slave endpoints: %w", err) } options := []struct { Key string Value any }{ {Key: "region_sync.enabled", Value: true}, {Key: "region_sync.region_id", Value: "ov"}, {Key: "region_sync.is_master", Value: false}, {Key: "region_sync.master_endpoint", Value: cnBaseURL}, {Key: "region_sync.slave_endpoints", Value: string(emptyEndpoints)}, {Key: "region_sync.sync_api_key", Value: apiKey}, } for _, option := range options { if err := SetOption(client, option.Key, option.Value); err != nil { return err } } return nil } func WaitOptionApplied(baseURL, apiKey string, expectedTotal int, timeout time.Duration) error { return WaitInternalMigrationReady(baseURL, apiKey, expectedTotal, timeout) } func WaitInternalMigrationReady(baseURL, apiKey string, expectedTotal int, timeout time.Duration) error { client, err := NewAPIClient(baseURL) if err != nil { return err } deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { var resp internalListUsersResponse err := client.DoJSONWithHeaders( http.MethodGet, "/api/internal/migration/users?page=0&page_size=1", nil, http.StatusOK, map[string]string{ "X-Sync-API-Key": apiKey, "X-Sync-Node": "e2e-ready-check", }, &resp, ) if err == nil && resp.Success && resp.Total == expectedTotal { return nil } time.Sleep(250 * time.Millisecond) } return fmt.Errorf("internal migration API at %s not ready within %s", baseURL, timeout) }