Co-Authored-By: Codex <noreply@anthropic.com>master
| @@ -0,0 +1,35 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetChannelAffinityUsageCacheStatsRequiresRuleAndFingerprint(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.GET("/stats", GetChannelAffinityUsageCacheStats) | |||
| for _, target := range []string{"/stats?key_fp=abc", "/stats?rule_name=rule"} { | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, target, nil)) | |||
| require.Equal(t, http.StatusBadRequest, response.Code) | |||
| require.Contains(t, response.Body.String(), `"success":false`) | |||
| } | |||
| } | |||
| func TestClearChannelAffinityCacheRequiresSelector(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.DELETE("/cache", ClearChannelAffinityCache) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodDelete, "/cache", nil)) | |||
| require.Equal(t, http.StatusBadRequest, response.Code) | |||
| require.Contains(t, response.Body.String(), `"success":false`) | |||
| } | |||
| @@ -0,0 +1,25 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetAuthHeadersUseProviderSpecificNames(t *testing.T) { | |||
| openAI := GetAuthHeader("openai-key") | |||
| claude := GetClaudeAuthHeader("claude-key") | |||
| require.Equal(t, "Bearer openai-key", openAI.Get("Authorization")) | |||
| require.Empty(t, openAI.Get("x-api-key")) | |||
| require.Equal(t, "claude-key", claude.Get("x-api-key")) | |||
| require.Equal(t, "2023-06-01", claude.Get("anthropic-version")) | |||
| require.Empty(t, claude.Get("Authorization")) | |||
| } | |||
| func TestGetResponseBodyRejectsInvalidMethodURLBeforeNetwork(t *testing.T) { | |||
| _, err := GetResponseBody("GET", "://invalid", nil, http.Header{}) | |||
| require.Error(t, err) | |||
| } | |||
| @@ -0,0 +1,45 @@ | |||
| package controller | |||
| import ( | |||
| "net/http/httptest" | |||
| "testing" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/pkg/ionet" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestMapIoNetDeploymentNormalizesStatusAndRemainingTime(t *testing.T) { | |||
| createdAt := time.Date(2026, 7, 25, 12, 0, 0, 0, time.UTC) | |||
| mapped := mapIoNetDeployment(ionet.Deployment{ | |||
| ID: "deployment-1", Name: "Model Host", Status: "RUNNING", CreatedAt: createdAt, | |||
| BrandName: "NVIDIA", HardwareName: "H100", HardwareQuantity: 2, ComputeMinutesRemaining: 125, | |||
| }) | |||
| require.Equal(t, "running", mapped["status"]) | |||
| require.Equal(t, "2 hour 5 minutes", mapped["time_remaining"]) | |||
| require.Equal(t, "NVIDIA H100 x2", mapped["hardware_info"]) | |||
| require.EqualValues(t, createdAt.Unix(), mapped["created_at"]) | |||
| } | |||
| func TestComputeStatusCountsIncludesKnownAndUnknownStatuses(t *testing.T) { | |||
| counts := computeStatusCounts(4, []ionet.Deployment{ | |||
| {Status: "RUNNING"}, {Status: "running"}, {Status: "custom"}, | |||
| }) | |||
| require.EqualValues(t, 4, counts["all"]) | |||
| require.EqualValues(t, 2, counts["running"]) | |||
| require.EqualValues(t, 0, counts["failed"]) | |||
| require.EqualValues(t, 1, counts["custom"]) | |||
| } | |||
| func TestRequireDeploymentAndContainerIDsRejectBlankParameters(t *testing.T) { | |||
| context, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| context.Params = gin.Params{{Key: "id", Value: " "}, {Key: "container_id", Value: ""}} | |||
| _, ok := requireDeploymentID(context) | |||
| require.False(t, ok) | |||
| _, ok = requireContainerID(context) | |||
| require.False(t, ok) | |||
| } | |||
| @@ -0,0 +1,18 @@ | |||
| package controller | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestBoolToString(t *testing.T) { | |||
| require.Equal(t, "true", boolToString(true)) | |||
| require.Equal(t, "false", boolToString(false)) | |||
| } | |||
| func TestGetLegalContentUsesEnglishOnlyForExactEnglishLanguage(t *testing.T) { | |||
| require.Equal(t, "English", getLegalContent("Chinese", "English", "en")) | |||
| require.Equal(t, "Chinese", getLegalContent("Chinese", "English", "zh")) | |||
| require.Equal(t, "Chinese", getLegalContent("Chinese", "English", "en-US")) | |||
| } | |||
| @@ -0,0 +1,40 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetOptionsOmitsSensitiveSuffixes(t *testing.T) { | |||
| common.OptionMapRWMutex.Lock() | |||
| oldOptions := common.OptionMap | |||
| common.OptionMap = map[string]string{ | |||
| "PublicOption": "visible", "AccessToken": "secret", "WebhookSecret": "secret", | |||
| "ProviderKey": "secret", "lowercase_secret": "secret", "service_api_key": "secret", | |||
| } | |||
| common.OptionMapRWMutex.Unlock() | |||
| t.Cleanup(func() { | |||
| common.OptionMapRWMutex.Lock() | |||
| common.OptionMap = oldOptions | |||
| common.OptionMapRWMutex.Unlock() | |||
| }) | |||
| router := gin.New() | |||
| router.GET("/", GetOptions) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/", nil)) | |||
| require.Equal(t, http.StatusOK, response.Code) | |||
| body := response.Body.String() | |||
| require.Contains(t, body, "PublicOption") | |||
| require.NotContains(t, body, "AccessToken") | |||
| require.NotContains(t, body, "WebhookSecret") | |||
| require.NotContains(t, body, "ProviderKey") | |||
| require.NotContains(t, body, "lowercase_secret") | |||
| require.NotContains(t, body, "service_api_key") | |||
| } | |||
| @@ -0,0 +1,57 @@ | |||
| package controller | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestStripeWebhookAvailabilityRequiresAllTopUpCredentials(t *testing.T) { | |||
| oldAPI, oldWebhook, oldPrice := setting.StripeApiSecret, setting.StripeWebhookSecret, setting.StripePriceId | |||
| t.Cleanup(func() { | |||
| setting.StripeApiSecret, setting.StripeWebhookSecret, setting.StripePriceId = oldAPI, oldWebhook, oldPrice | |||
| }) | |||
| setting.StripeApiSecret, setting.StripeWebhookSecret, setting.StripePriceId = "api", "webhook", "price" | |||
| require.True(t, isStripeTopUpEnabled()) | |||
| require.True(t, isStripeWebhookEnabled()) | |||
| setting.StripeWebhookSecret = " " | |||
| require.False(t, isStripeTopUpEnabled()) | |||
| require.False(t, isStripeWebhookEnabled()) | |||
| } | |||
| func TestCreemWebhookAvailabilityRequiresProductsAndWebhookSecret(t *testing.T) { | |||
| oldAPI, oldProducts, oldWebhook := setting.CreemApiKey, setting.CreemProducts, setting.CreemWebhookSecret | |||
| t.Cleanup(func() { | |||
| setting.CreemApiKey, setting.CreemProducts, setting.CreemWebhookSecret = oldAPI, oldProducts, oldWebhook | |||
| }) | |||
| setting.CreemApiKey, setting.CreemProducts, setting.CreemWebhookSecret = "api", "[]", "webhook" | |||
| require.False(t, isCreemTopUpEnabled()) | |||
| require.False(t, isCreemWebhookEnabled()) | |||
| setting.CreemProducts = `[{"id":"product"}]` | |||
| require.True(t, isCreemTopUpEnabled()) | |||
| require.True(t, isCreemWebhookEnabled()) | |||
| setting.CreemWebhookSecret = "" | |||
| require.False(t, isCreemWebhookEnabled()) | |||
| } | |||
| func TestEpayWebhookAvailabilityRequiresPaymentMethod(t *testing.T) { | |||
| oldAddress, oldID, oldKey, oldMethods := operation_setting.PayAddress, operation_setting.EpayId, operation_setting.EpayKey, operation_setting.PayMethods | |||
| t.Cleanup(func() { | |||
| operation_setting.PayAddress, operation_setting.EpayId, operation_setting.EpayKey, operation_setting.PayMethods = oldAddress, oldID, oldKey, oldMethods | |||
| }) | |||
| operation_setting.PayAddress, operation_setting.EpayId, operation_setting.EpayKey = "https://pay.example", "merchant", "secret" | |||
| operation_setting.PayMethods = nil | |||
| require.True(t, isEpayTopUpEnabled()) | |||
| require.False(t, isEpayWebhookEnabled()) | |||
| operation_setting.PayMethods = []map[string]string{{"name": "alipay"}} | |||
| require.True(t, isEpayWebhookEnabled()) | |||
| } | |||
| @@ -0,0 +1,31 @@ | |||
| package controller | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetRatioConfigRejectsRequestWhenExposureIsDisabled(t *testing.T) { | |||
| ratio_setting.SetExposeRatioEnabled(false) | |||
| t.Cleanup(func() { ratio_setting.SetExposeRatioEnabled(false) }) | |||
| context, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| GetRatioConfig(context) | |||
| require.Equal(t, http.StatusForbidden, context.Writer.Status()) | |||
| } | |||
| func TestGetRatioConfigReturnsDataWhenExposureIsEnabled(t *testing.T) { | |||
| ratio_setting.SetExposeRatioEnabled(true) | |||
| t.Cleanup(func() { ratio_setting.SetExposeRatioEnabled(false) }) | |||
| context, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| GetRatioConfig(context) | |||
| require.Equal(t, http.StatusOK, context.Writer.Status()) | |||
| } | |||
| @@ -0,0 +1,209 @@ | |||
| package controller | |||
| import ( | |||
| "fmt" | |||
| "net/http" | |||
| "strconv" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/QuantumNous/new-api/service" | |||
| "github.com/gin-gonic/gin" | |||
| "gorm.io/gorm" | |||
| ) | |||
| type AdminVideoChannelBinding struct { | |||
| Group string `json:"group"` | |||
| Family string `json:"family"` | |||
| ChannelID int `json:"channel_id"` | |||
| } | |||
| type adminVideoChannelBindingRequest struct { | |||
| Bindings []AdminVideoChannelBinding `json:"bindings"` | |||
| } | |||
| type validatedAdminVideoChannelBinding struct { | |||
| AdminVideoChannelBinding | |||
| ChannelType int | |||
| } | |||
| type AdminVideoChannelCandidate struct { | |||
| ID int `json:"id"` | |||
| Name string `json:"name"` | |||
| Type int `json:"type"` | |||
| } | |||
| type adminVideoChannelBindingRow struct { | |||
| Group string `json:"group"` | |||
| ChannelID int `json:"channel_id"` | |||
| ChannelName string `json:"channel_name"` | |||
| ChannelType int `json:"channel_type"` | |||
| Candidates []AdminVideoChannelCandidate `json:"candidates"` | |||
| } | |||
| type adminVideoChannelBindingFamily struct { | |||
| Key string `json:"key"` | |||
| Name string `json:"name"` | |||
| Bindings []adminVideoChannelBindingRow `json:"bindings"` | |||
| } | |||
| func videoAssetFamilyFromString(value string) (service.VideoAssetFamily, bool) { | |||
| family := service.VideoAssetFamily(strings.TrimSpace(value)) | |||
| for _, candidate := range service.VideoAssetFamilies() { | |||
| if family == candidate { | |||
| return family, true | |||
| } | |||
| } | |||
| return "", false | |||
| } | |||
| func GetUserVideoChannelBindings(c *gin.Context) { | |||
| userID, err := strconv.Atoi(c.Param("id")) | |||
| if err != nil || userID <= 0 { | |||
| common.ApiErrorMsg(c, "invalid user id") | |||
| return | |||
| } | |||
| if _, err = model.GetUserById(userID, false); err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| groups, err := model.GetUserConcreteTokenGroups(userID) | |||
| if err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| visibleGroupSet := make(map[string]struct{}, len(groups)) | |||
| families := make([]adminVideoChannelBindingFamily, 0, len(service.VideoAssetFamilies())) | |||
| for _, family := range service.VideoAssetFamilies() { | |||
| rows := make([]adminVideoChannelBindingRow, 0, len(groups)) | |||
| for _, group := range groups { | |||
| candidates, err := service.GetVideoAssetChannelCandidates(group, family) | |||
| if err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| if len(candidates) == 0 { | |||
| continue | |||
| } | |||
| visibleGroupSet[group] = struct{}{} | |||
| row := adminVideoChannelBindingRow{Group: group, Candidates: make([]AdminVideoChannelCandidate, 0, len(candidates))} | |||
| for _, channel := range candidates { | |||
| row.Candidates = append(row.Candidates, AdminVideoChannelCandidate{ID: channel.Id, Name: channel.Name, Type: channel.Type}) | |||
| } | |||
| bindings, err := model.GetUserAssetChannelsByTypes(userID, service.VideoAssetChannelTypesForFamily(family), group) | |||
| if err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| for _, binding := range bindings { | |||
| for _, candidate := range candidates { | |||
| if candidate.Id == binding.ChannelId { | |||
| row.ChannelID = candidate.Id | |||
| row.ChannelName = candidate.Name | |||
| row.ChannelType = candidate.Type | |||
| break | |||
| } | |||
| } | |||
| if row.ChannelID != 0 { | |||
| break | |||
| } | |||
| } | |||
| rows = append(rows, row) | |||
| } | |||
| families = append(families, adminVideoChannelBindingFamily{Key: string(family), Name: strings.ToUpper(string(family[:1])) + string(family[1:]), Bindings: rows}) | |||
| } | |||
| visibleGroups := make([]string, 0, len(visibleGroupSet)) | |||
| for _, group := range groups { | |||
| if _, ok := visibleGroupSet[group]; ok { | |||
| visibleGroups = append(visibleGroups, group) | |||
| } | |||
| } | |||
| c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"groups": visibleGroups, "families": families}}) | |||
| } | |||
| func SetUserVideoChannelBindings(c *gin.Context) { | |||
| userID, err := strconv.Atoi(c.Param("id")) | |||
| if err != nil || userID <= 0 { | |||
| common.ApiErrorMsg(c, "invalid user id") | |||
| return | |||
| } | |||
| if _, err = model.GetUserById(userID, false); err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| var request adminVideoChannelBindingRequest | |||
| if err := c.ShouldBindJSON(&request); err != nil { | |||
| common.ApiErrorMsg(c, err.Error()) | |||
| return | |||
| } | |||
| groups, err := model.GetUserConcreteTokenGroups(userID) | |||
| if err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| allowedGroups := make(map[string]struct{}, len(groups)) | |||
| for _, group := range groups { | |||
| allowedGroups[group] = struct{}{} | |||
| } | |||
| requested := make(map[string]validatedAdminVideoChannelBinding, len(request.Bindings)) | |||
| for _, binding := range request.Bindings { | |||
| binding.Group = strings.TrimSpace(binding.Group) | |||
| family, ok := videoAssetFamilyFromString(binding.Family) | |||
| if !ok || binding.Group == "" || binding.Group == "auto" || binding.ChannelID <= 0 { | |||
| common.ApiErrorMsg(c, "invalid video channel binding") | |||
| return | |||
| } | |||
| if _, ok := allowedGroups[binding.Group]; !ok { | |||
| common.ApiErrorMsg(c, "token group is not available for this user") | |||
| return | |||
| } | |||
| key := binding.Group + "\x00" + string(family) | |||
| if _, exists := requested[key]; exists { | |||
| common.ApiErrorMsg(c, "duplicate video channel binding") | |||
| return | |||
| } | |||
| candidates, err := service.GetVideoAssetChannelCandidates(binding.Group, family) | |||
| if err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| channelType := 0 | |||
| for _, candidate := range candidates { | |||
| if candidate.Id == binding.ChannelID { | |||
| channelType = candidate.Type | |||
| break | |||
| } | |||
| } | |||
| if channelType == 0 { | |||
| common.ApiErrorMsg(c, "channel is not available for this video family and token group") | |||
| return | |||
| } | |||
| binding.Family = string(family) | |||
| requested[key] = validatedAdminVideoChannelBinding{AdminVideoChannelBinding: binding, ChannelType: channelType} | |||
| } | |||
| if err := model.DB.Transaction(func(tx *gorm.DB) error { | |||
| for _, group := range groups { | |||
| for _, family := range service.VideoAssetFamilies() { | |||
| key := group + "\x00" + string(family) | |||
| binding, exists := requested[key] | |||
| channelTypes := service.VideoAssetChannelTypesForFamily(family) | |||
| if err := model.DeleteUserAssetChannelsByTypesWithTx(tx, userID, channelTypes, group); err != nil { | |||
| return err | |||
| } | |||
| if !exists { | |||
| continue | |||
| } | |||
| if err := model.BindUserAssetChannelWithTx(tx, userID, binding.ChannelType, group, binding.ChannelID); err != nil { | |||
| return err | |||
| } | |||
| } | |||
| } | |||
| return nil | |||
| }); err != nil { | |||
| common.ApiError(c, err) | |||
| return | |||
| } | |||
| model.RecordLog(userID, model.LogTypeManage, fmt.Sprintf("updated video channel bindings for user %d", userID)) | |||
| c.JSON(http.StatusOK, gin.H{"success": true, "message": ""}) | |||
| } | |||
| @@ -0,0 +1,243 @@ | |||
| package controller | |||
| import ( | |||
| "bytes" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "os" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func setupUserVideoChannelBindingDB(t *testing.T) *gorm.DB { | |||
| t.Helper() | |||
| originalDB := model.DB | |||
| originalLogDB := model.LOG_DB | |||
| originalCache := common.MemoryCacheEnabled | |||
| originalRedisEnabled := common.RedisEnabled | |||
| originalSQLitePath := common.SQLitePath | |||
| originalIsMasterNode := common.IsMasterNode | |||
| originalUsingSQLite := common.UsingSQLite | |||
| originalUsingMySQL := common.UsingMySQL | |||
| originalUsingPostgreSQL := common.UsingPostgreSQL | |||
| originalSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") | |||
| common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" | |||
| common.MemoryCacheEnabled = false | |||
| common.RedisEnabled = false | |||
| common.IsMasterNode = false | |||
| common.UsingSQLite = false | |||
| common.UsingMySQL = false | |||
| common.UsingPostgreSQL = false | |||
| require.NoError(t, os.Setenv("SQL_DSN", "local")) | |||
| require.NoError(t, model.InitDB()) | |||
| db := model.DB | |||
| model.LOG_DB = db | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.Log{})) | |||
| t.Cleanup(func() { | |||
| model.DB = originalDB | |||
| model.LOG_DB = originalLogDB | |||
| common.MemoryCacheEnabled = originalCache | |||
| common.RedisEnabled = originalRedisEnabled | |||
| common.SQLitePath = originalSQLitePath | |||
| common.IsMasterNode = originalIsMasterNode | |||
| common.UsingSQLite = originalUsingSQLite | |||
| common.UsingMySQL = originalUsingMySQL | |||
| common.UsingPostgreSQL = originalUsingPostgreSQL | |||
| if hadSQLDSN { | |||
| _ = os.Setenv("SQL_DSN", originalSQLDSN) | |||
| } else { | |||
| _ = os.Unsetenv("SQL_DSN") | |||
| } | |||
| require.NoError(t, sqlDB.Close()) | |||
| }) | |||
| return db | |||
| } | |||
| func setupUserVideoChannelBindingRouter() *gin.Engine { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.GET("/api/user/:id/video-channel-bindings", GetUserVideoChannelBindings) | |||
| router.PUT("/api/user/:id/video-channel-bindings", SetUserVideoChannelBindings) | |||
| return router | |||
| } | |||
| func createUserVideoBindingChannel(t *testing.T, db *gorm.DB, id, channelType int, group string, status int) { | |||
| t.Helper() | |||
| priority := int64(id) | |||
| weight := uint(1) | |||
| autoBan := 1 | |||
| require.NoError(t, db.Create(&model.Channel{Id: id, Type: channelType, Key: "channel-key", Status: status, Name: "channel", Group: group, Models: "model", Priority: &priority, Weight: &weight, AutoBan: &autoBan}).Error) | |||
| } | |||
| func putUserVideoChannelBindings(t *testing.T, router *gin.Engine, body string) *httptest.ResponseRecorder { | |||
| t.Helper() | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, "/api/user/10/video-channel-bindings", bytes.NewBufferString(body))) | |||
| return w | |||
| } | |||
| func TestAdminGetUserVideoChannelBindings(t *testing.T) { | |||
| db := setupUserVideoChannelBindingDB(t) | |||
| router := setupUserVideoChannelBindingRouter() | |||
| require.NoError(t, db.Create(&model.User{Id: 10, Username: "user"}).Error) | |||
| for i, group := range []string{"default", "vip", "auto"} { | |||
| require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "token-" + string(rune('0'+i)), Group: group}).Error) | |||
| } | |||
| createUserVideoBindingChannel(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", common.ChannelStatusEnabled) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/api/user/10/video-channel-bindings", nil)) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.Contains(t, w.Body.String(), `"default"`) | |||
| assert.NotContains(t, w.Body.String(), `"vip"`) | |||
| assert.NotContains(t, w.Body.String(), `"auto"`) | |||
| assert.Contains(t, w.Body.String(), `"channel_id":7`) | |||
| assert.NotContains(t, w.Body.String(), `channel-key`) | |||
| } | |||
| func TestAdminGetUserVideoChannelBindingsFiltersGroupsAndFamiliesWithoutCandidates(t *testing.T) { | |||
| db := setupUserVideoChannelBindingDB(t) | |||
| router := setupUserVideoChannelBindingRouter() | |||
| require.NoError(t, db.Create(&model.User{Id: 10, Username: "user"}).Error) | |||
| for i, group := range []string{"default", "vip", "test"} { | |||
| require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "token-" + group, Group: group}).Error) | |||
| } | |||
| createUserVideoBindingChannel(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 59, constant.ChannelTypeKlingAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "vip", common.ChannelStatusEnabled) | |||
| w := httptest.NewRecorder() | |||
| router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/api/user/10/video-channel-bindings", nil)) | |||
| var response struct { | |||
| Success bool `json:"success"` | |||
| Data struct { | |||
| Groups []string `json:"groups"` | |||
| Families []adminVideoChannelBindingFamily `json:"families"` | |||
| } `json:"data"` | |||
| } | |||
| require.NoError(t, common.Unmarshal(w.Body.Bytes(), &response)) | |||
| require.True(t, response.Success) | |||
| assert.Equal(t, []string{"default", "vip"}, response.Data.Groups) | |||
| familyGroups := map[string][]string{} | |||
| for _, family := range response.Data.Families { | |||
| for _, binding := range family.Bindings { | |||
| familyGroups[family.Key] = append(familyGroups[family.Key], binding.Group) | |||
| require.NotEmpty(t, binding.Candidates) | |||
| } | |||
| } | |||
| assert.Equal(t, []string{"default", "vip"}, familyGroups["seedance"]) | |||
| assert.Equal(t, []string{"default"}, familyGroups["kling"]) | |||
| } | |||
| func TestAdminSetUserVideoChannelBindings(t *testing.T) { | |||
| db := setupUserVideoChannelBindingDB(t) | |||
| router := setupUserVideoChannelBindingRouter() | |||
| require.NoError(t, db.Create(&model.User{Id: 10, Username: "user"}).Error) | |||
| require.NoError(t, db.Create(&model.Token{Id: 1, UserId: 10, Key: "token", Group: "default"}).Error) | |||
| createUserVideoBindingChannel(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 59, constant.ChannelTypeKlingAiping, "default", common.ChannelStatusEnabled) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59)) | |||
| w := putUserVideoChannelBindings(t, router, `{"bindings":[{"group":"default","family":"seedance","channel_id":16}]}`) | |||
| require.Equal(t, http.StatusOK, w.Code) | |||
| assert.Contains(t, w.Body.String(), `"success":true`) | |||
| oldBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, oldBinding) | |||
| newBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, newBinding) | |||
| assert.Equal(t, 16, newBinding.ChannelId) | |||
| klingBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default") | |||
| require.NoError(t, err) | |||
| assert.Nil(t, klingBinding) | |||
| } | |||
| func TestAdminSetUserVideoChannelBindingsRejectsInvalidRequestsWithoutChangingBindings(t *testing.T) { | |||
| db := setupUserVideoChannelBindingDB(t) | |||
| router := setupUserVideoChannelBindingRouter() | |||
| require.NoError(t, db.Create(&model.User{Id: 10, Username: "user"}).Error) | |||
| for i, group := range []string{"default", "auto"} { | |||
| require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "token-" + group, Group: group}).Error) | |||
| } | |||
| createUserVideoBindingChannel(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 59, constant.ChannelTypeKlingAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 60, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 61, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusAutoDisabled) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) | |||
| for _, body := range []string{ | |||
| `{"bindings":[{"group":"vip","family":"seedance","channel_id":16}]}`, | |||
| `{"bindings":[{"group":"auto","family":"seedance","channel_id":16}]}`, | |||
| `{"bindings":[{"group":"default","family":"kling","channel_id":16}]}`, | |||
| `{"bindings":[{"group":"default","family":"seedance","channel_id":60}]}`, | |||
| `{"bindings":[{"group":"default","family":"seedance","channel_id":61}]}`, | |||
| `{"bindings":[{"group":"default","family":"seedance","channel_id":16},{"group":"default","family":"seedance","channel_id":7}]}`, | |||
| } { | |||
| w := putUserVideoChannelBindings(t, router, body) | |||
| assert.Contains(t, w.Body.String(), `"success":false`) | |||
| binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, binding) | |||
| assert.Equal(t, 7, binding.ChannelId) | |||
| } | |||
| } | |||
| func TestAdminSetUserVideoChannelBindingsPreservesLegacyGroupBinding(t *testing.T) { | |||
| db := setupUserVideoChannelBindingDB(t) | |||
| router := setupUserVideoChannelBindingRouter() | |||
| require.NoError(t, db.Create(&model.User{Id: 10, Username: "user"}).Error) | |||
| require.NoError(t, db.Create(&model.Token{Id: 1, UserId: 10, Key: "token", Group: "default"}).Error) | |||
| createUserVideoBindingChannel(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusEnabled) | |||
| createUserVideoBindingChannel(t, db, 59, constant.ChannelTypeKlingAiping, "legacy", common.ChannelStatusEnabled) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "legacy", 59)) | |||
| w := putUserVideoChannelBindings(t, router, `{"bindings":[{"group":"default","family":"seedance","channel_id":7}]}`) | |||
| assert.Contains(t, w.Body.String(), `"success":true`) | |||
| legacyBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "legacy") | |||
| require.NoError(t, err) | |||
| require.NotNil(t, legacyBinding) | |||
| assert.Equal(t, 59, legacyBinding.ChannelId) | |||
| } | |||
| func TestAdminSetUserVideoChannelBindingsClearsHiddenCurrentGroupBinding(t *testing.T) { | |||
| db := setupUserVideoChannelBindingDB(t) | |||
| router := setupUserVideoChannelBindingRouter() | |||
| require.NoError(t, db.Create(&model.User{Id: 10, Username: "user"}).Error) | |||
| for i, group := range []string{"default", "test"} { | |||
| require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "token-" + group, Group: group}).Error) | |||
| } | |||
| createUserVideoBindingChannel(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", common.ChannelStatusEnabled) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "test", 99)) | |||
| w := putUserVideoChannelBindings(t, router, `{"bindings":[]}`) | |||
| assert.Contains(t, w.Body.String(), `"success":true`) | |||
| for _, group := range []string{"default", "test"} { | |||
| binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, group) | |||
| require.NoError(t, err) | |||
| assert.Nil(t, binding) | |||
| } | |||
| } | |||
| @@ -0,0 +1,40 @@ | |||
| package controller | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestExtractGeminiVideoURLSupportsTopLevelAndNestedShapes(t *testing.T) { | |||
| require.Equal(t, "https://video.example/top", extractGeminiVideoURLFromMap(map[string]any{"uri": "https://video.example/top"})) | |||
| nested := map[string]any{ | |||
| "response": map[string]any{ | |||
| "generateVideoResponse": map[string]any{ | |||
| "generatedSamples": []any{map[string]any{"video": map[string]any{"uri": "https://video.example/generated"}}}, | |||
| }, | |||
| }, | |||
| } | |||
| require.Equal(t, "https://video.example/generated", extractGeminiVideoURLFromMap(nested)) | |||
| require.Equal(t, "", extractGeminiVideoURLFromMap(map[string]any{"response": map[string]any{}})) | |||
| } | |||
| func TestExtractGeminiVideoURLFromTaskDataRejectsInvalidJSON(t *testing.T) { | |||
| task := &model.Task{Data: []byte("not-json")} | |||
| require.Empty(t, extractGeminiVideoURLFromTaskData(task)) | |||
| payload, err := common.Marshal(map[string]any{"response": map[string]any{"video": "https://video.example/nested"}}) | |||
| require.NoError(t, err) | |||
| task.Data = payload | |||
| require.Equal(t, "https://video.example/nested", extractGeminiVideoURLFromTaskData(task)) | |||
| } | |||
| func TestEnsureAPIKeyAppendsOnlyWhenMissing(t *testing.T) { | |||
| require.Equal(t, "https://video.example/file?key=abc", ensureAPIKey("https://video.example/file", "abc")) | |||
| require.Equal(t, "https://video.example/file?alt=media&key=abc", ensureAPIKey("https://video.example/file?alt=media", "abc")) | |||
| require.Equal(t, "https://video.example/file?key=existing", ensureAPIKey("https://video.example/file?key=existing", "abc")) | |||
| require.Equal(t, "https://video.example/file", ensureAPIKey("https://video.example/file", "")) | |||
| } | |||
| @@ -0,0 +1,30 @@ | |||
| # 管理员视频渠道绑定验收记录 | |||
| 日期:2026-07-24 | |||
| ## 自动化验证 | |||
| - `go test ./controller -run '^TestAdmin(Set|Get)UserVideoChannelBindings' -count=1`:通过。 | |||
| - GET 只返回用户的具体 Token 分组,不含 `auto`,且不泄露渠道密钥。 | |||
| - PUT 可切换 Seedance 渠道;省略的家族绑定被清空。 | |||
| - 拒绝非用户分组、`auto`、错误家族、跨分组、禁用渠道和重复绑定;失败不改变原绑定。 | |||
| - 仅覆盖当前 Token 分组,历史 `legacy` 分组的绑定保持不变。 | |||
| - `go test ./...`:通过。 | |||
| - `cd web; bun run build`:通过。保留项目既有的循环分包与大包告警。 | |||
| ## 手工验收步骤 | |||
| 1. 用管理员身份打开“用户管理”,编辑一个同时拥有 `default`、`vip` 和 `auto` Token 的用户。 | |||
| 2. 在“视频渠道绑定”卡片中点击“配置视频渠道”。确认只出现 `default`、`vip`,每个分组显示 Seedance 和 Kling。 | |||
| 3. 为 `default / Seedance` 选择一个候选渠道并保存;重新打开弹窗,确认选择仍在。 | |||
| 4. 清空 `default / Seedance` 并保存;重新打开弹窗,确认显示“未绑定”。 | |||
| 5. 使用该用户 `default` Token 发起对应 Seedance 请求:有绑定时命中指定渠道;清空后恢复既有自动选择。 | |||
| 6. 准备一个无任何 Seedance/Kling 候选渠道的 Token group,确认整个 group 不显示。 | |||
| 7. 准备一个只有 Seedance 候选渠道的 Token group,确认显示该 group 和 Seedance,但不显示 Kling。 | |||
| 8. 当所有具体 Token group 均无候选渠道时,确认弹窗显示“该用户当前 Token 分组没有可用的视频渠道”。 | |||
| 记录请求时只保存用户 ID、Token group、模型、渠道 ID 和 HTTP 状态,禁止保存 Token、渠道 Key、AK/SK 或签名 URL。 | |||
| ## 已知验证限制 | |||
| `bun run i18n:extract`、`bun run i18n:sync` 与 `bun run i18n:lint` 当前均因依赖版本不匹配失败:`react-i18next` 请求 `i18next.keyFromSelector`,但安装的 `i18next` 未导出该符号。新增中英文键已手动同步,生产构建已验证通过。 | |||
| @@ -0,0 +1,64 @@ | |||
| package middleware | |||
| import ( | |||
| "bytes" | |||
| "compress/gzip" | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/andybalholm/brotli" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestDecompressRequestMiddlewareExposesGzipAndBrotliPayload(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| for _, encoding := range []string{"gzip", "br"} { | |||
| t.Run(encoding, func(t *testing.T) { | |||
| var compressed bytes.Buffer | |||
| var writer io.WriteCloser | |||
| if encoding == "gzip" { | |||
| writer = gzip.NewWriter(&compressed) | |||
| } else { | |||
| writer = brotli.NewWriter(&compressed) | |||
| } | |||
| _, err := writer.Write([]byte(`{"message":"hello"}`)) | |||
| require.NoError(t, err) | |||
| require.NoError(t, writer.Close()) | |||
| router := gin.New() | |||
| router.Use(DecompressRequestMiddleware()) | |||
| router.POST("/", func(c *gin.Context) { | |||
| body, err := io.ReadAll(c.Request.Body) | |||
| require.NoError(t, err) | |||
| require.Equal(t, `{"message":"hello"}`, string(body)) | |||
| require.Empty(t, c.GetHeader("Content-Encoding")) | |||
| c.Status(http.StatusNoContent) | |||
| }) | |||
| request := httptest.NewRequest(http.MethodPost, "/", &compressed) | |||
| request.Header.Set("Content-Encoding", encoding) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, request) | |||
| require.Equal(t, http.StatusNoContent, response.Code) | |||
| }) | |||
| } | |||
| } | |||
| func TestDecompressRequestMiddlewareRejectsMalformedGzip(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| reachedHandler := false | |||
| router := gin.New() | |||
| router.Use(DecompressRequestMiddleware()) | |||
| router.POST("/", func(c *gin.Context) { reachedHandler = true }) | |||
| request := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString("not-gzip")) | |||
| request.Header.Set("Content-Encoding", "gzip") | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, request) | |||
| require.Equal(t, http.StatusBadRequest, response.Code) | |||
| require.False(t, reachedHandler) | |||
| } | |||
| @@ -0,0 +1,44 @@ | |||
| package middleware | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestDisableCacheSetsNoCacheHeaders(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(DisableCache()) | |||
| router.GET("/", func(c *gin.Context) { c.Status(http.StatusNoContent) }) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/", nil)) | |||
| require.Equal(t, http.StatusNoContent, response.Code) | |||
| require.Equal(t, "no-store, no-cache, must-revalidate, private, max-age=0", response.Header().Get("Cache-Control")) | |||
| require.Equal(t, "no-cache", response.Header().Get("Pragma")) | |||
| require.Equal(t, "0", response.Header().Get("Expires")) | |||
| } | |||
| func TestRequestIdPropagatesSameIDToContextAndResponse(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(RequestId()) | |||
| router.GET("/", func(c *gin.Context) { | |||
| id := c.GetString(common.RequestIdKey) | |||
| require.NotEmpty(t, id) | |||
| require.Equal(t, id, c.Request.Context().Value(common.RequestIdKey)) | |||
| c.Status(http.StatusNoContent) | |||
| }) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/", nil)) | |||
| require.Equal(t, http.StatusNoContent, response.Code) | |||
| require.NotEmpty(t, response.Header().Get(common.RequestIdKey)) | |||
| } | |||
| @@ -0,0 +1,41 @@ | |||
| package middleware | |||
| import ( | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestRelayPanicRecoverConvertsPanicToServerError(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(RelayPanicRecover()) | |||
| router.GET("/", func(c *gin.Context) { panic("upstream exploded") }) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/", nil)) | |||
| require.Equal(t, http.StatusInternalServerError, response.Code) | |||
| require.Contains(t, response.Body.String(), "new_api_panic") | |||
| require.Contains(t, response.Body.String(), "upstream exploded") | |||
| } | |||
| func TestTurnstileCheckAllowsRequestWhenFeatureIsDisabled(t *testing.T) { | |||
| oldEnabled := common.TurnstileCheckEnabled | |||
| common.TurnstileCheckEnabled = false | |||
| t.Cleanup(func() { common.TurnstileCheckEnabled = oldEnabled }) | |||
| gin.SetMode(gin.TestMode) | |||
| router := gin.New() | |||
| router.Use(TurnstileCheck()) | |||
| router.GET("/", func(c *gin.Context) { c.Status(http.StatusNoContent) }) | |||
| response := httptest.NewRecorder() | |||
| router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/", nil)) | |||
| require.Equal(t, http.StatusNoContent, response.Code) | |||
| } | |||
| @@ -0,0 +1,44 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func validCustomOAuthProvider() *CustomOAuthProvider { | |||
| return &CustomOAuthProvider{ | |||
| Name: "Example OAuth", | |||
| Slug: "Example-Provider", | |||
| ClientId: "client-id", | |||
| AuthorizationEndpoint: "https://id.example/authorize", | |||
| TokenEndpoint: "https://id.example/token", | |||
| UserInfoEndpoint: "https://id.example/userinfo", | |||
| } | |||
| } | |||
| func TestValidateCustomOAuthProviderNormalizesSlugAndAppliesDefaults(t *testing.T) { | |||
| provider := validCustomOAuthProvider() | |||
| err := validateCustomOAuthProvider(provider) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "example-provider", provider.Slug) | |||
| require.Equal(t, "sub", provider.UserIdField) | |||
| require.Equal(t, "preferred_username", provider.UsernameField) | |||
| require.Equal(t, "openid profile email", provider.Scopes) | |||
| } | |||
| func TestValidateCustomOAuthProviderRejectsInvalidSlugAndPolicy(t *testing.T) { | |||
| invalidSlug := validCustomOAuthProvider() | |||
| invalidSlug.Slug = "bad_slug" | |||
| require.ErrorContains(t, validateCustomOAuthProvider(invalidSlug), "slug") | |||
| unsupportedOp := validCustomOAuthProvider() | |||
| unsupportedOp.AccessPolicy = `{"conditions":[{"field":"role","op":"matches","value":"admin"}]}` | |||
| require.ErrorContains(t, validateCustomOAuthProvider(unsupportedOp), "unsupported") | |||
| nonArrayMembership := validCustomOAuthProvider() | |||
| nonArrayMembership.AccessPolicy = `{"conditions":[{"field":"role","op":"in","value":"admin"}]}` | |||
| require.ErrorContains(t, validateCustomOAuthProvider(nonArrayMembership), "must be an array") | |||
| } | |||
| @@ -0,0 +1,45 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestUpdateOptionMapSynchronizesBooleanAndNumericRuntimeSettings(t *testing.T) { | |||
| oldTaskEnabled := common.TaskEnabled | |||
| oldPrice := operation_setting.Price | |||
| oldOptionMap := common.OptionMap | |||
| common.OptionMap = make(map[string]string) | |||
| t.Cleanup(func() { | |||
| common.TaskEnabled = oldTaskEnabled | |||
| operation_setting.Price = oldPrice | |||
| common.OptionMap = oldOptionMap | |||
| }) | |||
| require.NoError(t, updateOptionMap("TaskEnabled", "false")) | |||
| require.False(t, common.TaskEnabled) | |||
| require.Equal(t, "false", common.OptionMap["TaskEnabled"]) | |||
| require.NoError(t, updateOptionMap("Price", "2.5")) | |||
| require.Equal(t, 2.5, operation_setting.Price) | |||
| } | |||
| func TestUpdateOptionMapRejectsInvalidStructuredSettingsWithoutReplacingExistingValue(t *testing.T) { | |||
| oldOptionMap := common.OptionMap | |||
| common.OptionMap = make(map[string]string) | |||
| oldAutoGroups := setting.AutoGroups2JsonString() | |||
| t.Cleanup(func() { | |||
| common.OptionMap = oldOptionMap | |||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(oldAutoGroups)) | |||
| }) | |||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default"]`)) | |||
| err := updateOptionMap("AutoGroups", "not-json") | |||
| require.Error(t, err) | |||
| require.Equal(t, []string{"default"}, setting.GetAutoGroups()) | |||
| } | |||
| @@ -0,0 +1,42 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/go-webauthn/webauthn/protocol" | |||
| webauthn "github.com/go-webauthn/webauthn/webauthn" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestPasskeyCredentialTransportsRoundTripAndIgnoreMalformedJSON(t *testing.T) { | |||
| credential := &PasskeyCredential{} | |||
| credential.SetTransports([]protocol.AuthenticatorTransport{protocol.USB, protocol.Internal}) | |||
| require.Equal(t, []protocol.AuthenticatorTransport{protocol.USB, protocol.Internal}, credential.TransportList()) | |||
| credential.Transports = "not-json" | |||
| require.Nil(t, credential.TransportList()) | |||
| credential.SetTransports(nil) | |||
| require.Empty(t, credential.Transports) | |||
| } | |||
| func TestPasskeyCredentialConvertsWebAuthnFieldsWithoutLosingFlags(t *testing.T) { | |||
| webCredential := &webauthn.Credential{ | |||
| ID: []byte("credential"), PublicKey: []byte("public-key"), AttestationType: "none", | |||
| Transport: []protocol.AuthenticatorTransport{protocol.Internal}, | |||
| Flags: webauthn.CredentialFlags{UserPresent: true, UserVerified: true, BackupEligible: true}, | |||
| Authenticator: webauthn.Authenticator{AAGUID: []byte("aaguid"), SignCount: 8, Attachment: protocol.Platform}, | |||
| } | |||
| stored := NewPasskeyCredentialFromWebAuthn(9, webCredential) | |||
| require.NotNil(t, stored) | |||
| require.Equal(t, 9, stored.UserID) | |||
| require.True(t, stored.UserVerified) | |||
| require.EqualValues(t, 8, stored.SignCount) | |||
| roundTripped := stored.ToWebAuthnCredential() | |||
| require.Equal(t, webCredential.ID, roundTripped.ID) | |||
| require.Equal(t, webCredential.PublicKey, roundTripped.PublicKey) | |||
| require.True(t, roundTripped.Flags.UserVerified) | |||
| require.Equal(t, webCredential.Transport, roundTripped.Transport) | |||
| } | |||
| @@ -0,0 +1,41 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestJSONValueScanCopiesByteInputAndSupportsStringAndNil(t *testing.T) { | |||
| input := []byte(`{"items":["a"]}`) | |||
| var value JSONValue | |||
| require.NoError(t, value.Scan(input)) | |||
| input[2] = 'X' | |||
| require.Equal(t, `{"items":["a"]}`, string(value)) | |||
| require.NoError(t, value.Scan(`{"items":["b"]}`)) | |||
| require.Equal(t, `{"items":["b"]}`, string(value)) | |||
| require.NoError(t, value.Scan(nil)) | |||
| require.Nil(t, value) | |||
| } | |||
| func TestJSONValueDatabaseAndJSONMarshallingPreservesRawPayload(t *testing.T) { | |||
| value := JSONValue(`{"models":["gpt-5"]}`) | |||
| databaseValue, err := value.Value() | |||
| require.NoError(t, err) | |||
| require.Equal(t, []byte(`{"models":["gpt-5"]}`), databaseValue) | |||
| encoded, err := value.MarshalJSON() | |||
| require.NoError(t, err) | |||
| require.Equal(t, []byte(`{"models":["gpt-5"]}`), encoded) | |||
| var decoded JSONValue | |||
| require.NoError(t, decoded.UnmarshalJSON([]byte(`["a","b"]`))) | |||
| require.Equal(t, `["a","b"]`, string(decoded)) | |||
| var nilValue JSONValue | |||
| encoded, err = nilValue.MarshalJSON() | |||
| require.NoError(t, err) | |||
| require.Equal(t, []byte("null"), encoded) | |||
| } | |||
| @@ -0,0 +1,13 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetDefaultVendorIconReturnsKnownIconAndEmptyFallback(t *testing.T) { | |||
| require.Equal(t, "OpenAI", getDefaultVendorIcon("OpenAI")) | |||
| require.Equal(t, "Claude.Color", getDefaultVendorIcon("Anthropic")) | |||
| require.Equal(t, "", getDefaultVendorIcon("Unknown Vendor")) | |||
| } | |||
| @@ -3,6 +3,7 @@ package model | |||
| import ( | |||
| "errors" | |||
| "fmt" | |||
| "sort" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| @@ -27,7 +28,7 @@ type Token struct { | |||
| AllowIps *string `json:"allow_ips" gorm:"default:''"` | |||
| UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota | |||
| Group string `json:"group" gorm:"default:''"` | |||
| CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试,仅auto分组有效 | |||
| CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试,仅auto分组有效 | |||
| DeletedAt gorm.DeletedAt `gorm:"index"` | |||
| } | |||
| @@ -64,6 +65,24 @@ func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { | |||
| return tokens, err | |||
| } | |||
| // GetUserConcreteTokenGroups returns the groups an administrator can bind. | |||
| // The auto group resolves at request time and cannot be bound directly. | |||
| func GetUserConcreteTokenGroups(userId int) ([]string, error) { | |||
| var rawGroups []string | |||
| if err := DB.Model(&Token{}).Where("user_id = ?", userId).Distinct().Pluck("group", &rawGroups).Error; err != nil { | |||
| return nil, err | |||
| } | |||
| groups := make([]string, 0, len(rawGroups)) | |||
| for _, group := range rawGroups { | |||
| group = strings.TrimSpace(group) | |||
| if group != "" && group != "auto" { | |||
| groups = append(groups, group) | |||
| } | |||
| } | |||
| sort.Strings(groups) | |||
| return groups, nil | |||
| } | |||
| // sanitizeLikePattern 校验并清洗用户输入的 LIKE 搜索模式。 | |||
| // 规则: | |||
| // 1. 转义 ! 和 _(使用 ! 作为 ESCAPE 字符,兼容 MySQL/PostgreSQL/SQLite) | |||
| @@ -0,0 +1,38 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestSanitizeLikePatternEscapesLiteralCharactersAndRejectsExpensiveWildcards(t *testing.T) { | |||
| pattern, err := sanitizeLikePattern("api_key!v1") | |||
| require.NoError(t, err) | |||
| require.Equal(t, "api!_key!!v1", pattern) | |||
| pattern, err = sanitizeLikePattern("ab%cd") | |||
| require.NoError(t, err) | |||
| require.Equal(t, "ab%cd", pattern) | |||
| _, err = sanitizeLikePattern("%%") | |||
| require.Error(t, err) | |||
| _, err = sanitizeLikePattern("a%b%c%d") | |||
| require.Error(t, err) | |||
| _, err = sanitizeLikePattern("a%") | |||
| require.Error(t, err) | |||
| } | |||
| func TestTokenParsersNormalizeIPAndModelLimitFields(t *testing.T) { | |||
| allowIPs := " 127.0.0.1,\n 10.0.0.1 \n\n" | |||
| token := Token{AllowIps: &allowIPs, ModelLimits: "gpt-5,claude-sonnet"} | |||
| require.Equal(t, []string{"127.0.0.1", "10.0.0.1"}, token.GetIpLimits()) | |||
| require.Equal(t, []string{"gpt-5", "claude-sonnet"}, token.GetModelLimits()) | |||
| require.Equal(t, map[string]bool{"gpt-5": true, "claude-sonnet": true}, token.GetModelLimitsMap()) | |||
| token.AllowIps = nil | |||
| token.ModelLimits = "" | |||
| require.Empty(t, token.GetIpLimits()) | |||
| require.Empty(t, token.GetModelLimits()) | |||
| } | |||
| @@ -0,0 +1,21 @@ | |||
| package model | |||
| import ( | |||
| "testing" | |||
| "time" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestTwoFAIsLockedDependsOnFutureLockDeadline(t *testing.T) { | |||
| noDeadline := &TwoFA{} | |||
| require.False(t, noDeadline.IsLocked()) | |||
| future := time.Now().Add(time.Minute) | |||
| locked := &TwoFA{LockedUntil: &future} | |||
| require.True(t, locked.IsLocked()) | |||
| past := time.Now().Add(-time.Minute) | |||
| expired := &TwoFA{LockedUntil: &past} | |||
| require.False(t, expired.IsLocked()) | |||
| } | |||
| @@ -0,0 +1,42 @@ | |||
| package model | |||
| import ( | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestUserBaseGetSettingHandlesValidAndInvalidJSON(t *testing.T) { | |||
| user := UserBase{Setting: `{"language":"en","billing_preference":"subscription"}`} | |||
| setting := user.GetSetting() | |||
| require.Equal(t, "en", setting.Language) | |||
| require.Equal(t, "subscription", setting.BillingPreference) | |||
| invalid := UserBase{Setting: "{"} | |||
| require.Equal(t, "", invalid.GetSetting().Language) | |||
| } | |||
| func TestUserBaseWriteContextPopulatesRelayFields(t *testing.T) { | |||
| context, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| user := UserBase{ | |||
| Source: "oauth", Group: "vip", Quota: 123, Status: common.UserStatusEnabled, | |||
| Email: "user@example.com", Username: "user", Setting: `{"language":"en"}`, | |||
| } | |||
| user.WriteContext(context) | |||
| require.Equal(t, "vip", common.GetContextKeyString(context, constant.ContextKeyUserGroup)) | |||
| require.Equal(t, 123, common.GetContextKeyInt(context, constant.ContextKeyUserQuota)) | |||
| require.Equal(t, common.UserStatusEnabled, common.GetContextKeyInt(context, constant.ContextKeyUserStatus)) | |||
| require.Equal(t, "user@example.com", common.GetContextKeyString(context, constant.ContextKeyUserEmail)) | |||
| require.Equal(t, "user", common.GetContextKeyString(context, constant.ContextKeyUserName)) | |||
| require.Equal(t, "oauth", common.GetContextKeyString(context, constant.ContextKeyUserSource)) | |||
| } | |||
| func TestGetUserCacheKeyIsNamespacedByUserID(t *testing.T) { | |||
| require.Equal(t, "user:42", getUserCacheKey(42)) | |||
| } | |||
| @@ -0,0 +1,38 @@ | |||
| package model | |||
| import ( | |||
| "errors" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/stretchr/testify/require" | |||
| "gorm.io/gorm" | |||
| ) | |||
| func TestRecordExistClassifiesSuccessNotFoundAndDatabaseError(t *testing.T) { | |||
| exists, err := RecordExist(nil) | |||
| require.True(t, exists) | |||
| require.NoError(t, err) | |||
| exists, err = RecordExist(gorm.ErrRecordNotFound) | |||
| require.False(t, exists) | |||
| require.NoError(t, err) | |||
| dbErr := errors.New("database unavailable") | |||
| exists, err = RecordExist(dbErr) | |||
| require.False(t, exists) | |||
| require.ErrorIs(t, err, dbErr) | |||
| } | |||
| func TestShouldUpdateRedisRequiresRedisDatabaseSourceAndNoError(t *testing.T) { | |||
| oldRedisEnabled := common.RedisEnabled | |||
| t.Cleanup(func() { common.RedisEnabled = oldRedisEnabled }) | |||
| common.RedisEnabled = false | |||
| require.False(t, shouldUpdateRedis(true, nil)) | |||
| common.RedisEnabled = true | |||
| require.False(t, shouldUpdateRedis(false, nil)) | |||
| require.False(t, shouldUpdateRedis(true, errors.New("query failed"))) | |||
| require.True(t, shouldUpdateRedis(true, nil)) | |||
| } | |||
| @@ -136,6 +136,8 @@ func SetApiRouter(router *gin.Engine) { | |||
| adminRoute.GET("/:id/oauth/bindings", controller.GetUserOAuthBindingsByAdmin) | |||
| adminRoute.DELETE("/:id/oauth/bindings/:provider_id", controller.UnbindCustomOAuthByAdmin) | |||
| adminRoute.DELETE("/:id/bindings/:binding_type", controller.AdminClearUserBinding) | |||
| adminRoute.GET("/:id/video-channel-bindings", controller.GetUserVideoChannelBindings) | |||
| adminRoute.PUT("/:id/video-channel-bindings", controller.SetUserVideoChannelBindings) | |||
| adminRoute.GET("/:id", controller.GetUser) | |||
| adminRoute.POST("/", controller.CreateUser) | |||
| adminRoute.POST("/manage", controller.ManageUser) | |||
| @@ -0,0 +1,51 @@ | |||
| package service | |||
| import ( | |||
| "encoding/base64" | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestParseAudioUsesFormatSpecificSampleRates(t *testing.T) { | |||
| pcm := base64.StdEncoding.EncodeToString(make([]byte, 48_000)) | |||
| duration, err := parseAudio(pcm, "pcm16") | |||
| require.NoError(t, err) | |||
| require.Equal(t, 1.0, duration) | |||
| g711 := base64.StdEncoding.EncodeToString(make([]byte, 8_000)) | |||
| duration, err = parseAudio(g711, "g711_ulaw") | |||
| require.NoError(t, err) | |||
| require.Equal(t, 1.0, duration) | |||
| } | |||
| func TestParseAudioAndDecodeBase64AudioDataRejectInvalidEncoding(t *testing.T) { | |||
| _, err := parseAudio("not base64", "pcm16") | |||
| require.Error(t, err) | |||
| _, err = DecodeBase64AudioData("not base64") | |||
| require.Error(t, err) | |||
| } | |||
| func TestDecodeBase64AudioDataStripsDataURLPrefix(t *testing.T) { | |||
| decoded, err := DecodeBase64AudioData("data:audio/pcm;base64,AAE=") | |||
| require.NoError(t, err) | |||
| require.Equal(t, "AAE=", decoded) | |||
| } | |||
| func TestCountAudioTokensUsesConfiguredInputAndOutputRates(t *testing.T) { | |||
| oneSecondPCM := base64.StdEncoding.EncodeToString(make([]byte, 48_000)) | |||
| inputTokens, err := CountAudioTokenInput(oneSecondPCM, "pcm16") | |||
| require.NoError(t, err) | |||
| require.Equal(t, 27, inputTokens) | |||
| outputTokens, err := CountAudioTokenOutput(oneSecondPCM, "pcm16") | |||
| require.NoError(t, err) | |||
| require.Equal(t, 13, outputTokens) | |||
| inputTokens, err = CountAudioTokenInput("", "pcm16") | |||
| require.NoError(t, err) | |||
| require.Zero(t, inputTokens) | |||
| } | |||
| @@ -0,0 +1,35 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestChannelAffinityMatchersIgnoreInvalidPatternsAndMatchCaseInsensitively(t *testing.T) { | |||
| require.True(t, matchAnyRegexCached([]string{"[", `^gpt-[0-9]+$`}, "gpt-5")) | |||
| require.False(t, matchAnyRegexCached([]string{"["}, "gpt-5")) | |||
| require.True(t, matchAnyIncludeFold([]string{" Codex ", "other"}, "Mozilla CodexClient")) | |||
| require.False(t, matchAnyIncludeFold([]string{""}, "CodexClient")) | |||
| } | |||
| func TestChannelAffinityKeyHelpersProduceStableSafeValues(t *testing.T) { | |||
| rule := operation_setting.ChannelAffinityRule{Name: "rule", IncludeRuleName: true, IncludeUsingGroup: true} | |||
| require.Equal(t, "rule:vip:key", buildChannelAffinityCacheKeySuffix(rule, "vip", "key")) | |||
| require.Equal(t, "abcd...wxyz", buildChannelAffinityKeyHint("abcdefghijklmnopqrstuvwxwxyz")) | |||
| require.Len(t, affinityFingerprint("tenant-123"), 8) | |||
| require.Equal(t, "", channelAffinityUsageCacheEntryKey("", "vip", "fingerprint")) | |||
| require.Equal(t, "rule\n\nfingerprint", channelAffinityUsageCacheEntryKey("rule", "", "fingerprint")) | |||
| } | |||
| func TestChannelAffinityUsageHelpersUseFallbackFields(t *testing.T) { | |||
| usage := &dto.Usage{InputTokens: 3, OutputTokens: 5} | |||
| require.Equal(t, 3, usagePromptTokens(usage)) | |||
| require.Equal(t, 5, usageCompletionTokens(usage)) | |||
| require.Equal(t, 8, usageTotalTokens(usage)) | |||
| require.Equal(t, cacheTokenRateModeCachedOverPrompt, cachedTokenRateModeByRelayFormat(types.RelayFormatOpenAI)) | |||
| require.Equal(t, cacheTokenRateModeCachedOverPromptPlusCached, cachedTokenRateModeByRelayFormat(types.RelayFormatClaude)) | |||
| } | |||
| @@ -0,0 +1,67 @@ | |||
| package service | |||
| import ( | |||
| "errors" | |||
| "net/http" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestShouldDisableChannelHonorsFeatureFlagAndChannelErrors(t *testing.T) { | |||
| original := common.AutomaticDisableChannelEnabled | |||
| t.Cleanup(func() { common.AutomaticDisableChannelEnabled = original }) | |||
| channelErr := types.NewError(errors.New("no usable key"), types.ErrorCodeChannelNoAvailableKey) | |||
| common.AutomaticDisableChannelEnabled = false | |||
| require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, channelErr)) | |||
| common.AutomaticDisableChannelEnabled = true | |||
| require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, nil)) | |||
| require.True(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, channelErr)) | |||
| skipRetry := types.NewError(errors.New("retry later"), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry()) | |||
| require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, skipRetry)) | |||
| } | |||
| func TestShouldDisableChannelRecognizesStatusAndOpenAIErrorRules(t *testing.T) { | |||
| originalEnabled := common.AutomaticDisableChannelEnabled | |||
| originalRanges := operation_setting.AutomaticDisableStatusCodeRanges | |||
| t.Cleanup(func() { | |||
| common.AutomaticDisableChannelEnabled = originalEnabled | |||
| operation_setting.AutomaticDisableStatusCodeRanges = originalRanges | |||
| }) | |||
| common.AutomaticDisableChannelEnabled = true | |||
| operation_setting.AutomaticDisableStatusCodeRanges = []operation_setting.StatusCodeRange{{Start: http.StatusTooManyRequests, End: http.StatusTooManyRequests}} | |||
| byStatus := types.NewOpenAIError(errors.New("rate limited"), types.ErrorCodeBadResponse, http.StatusTooManyRequests) | |||
| require.True(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, byStatus)) | |||
| forbidden := types.NewOpenAIError(errors.New("forbidden"), types.ErrorCodeBadResponse, http.StatusForbidden) | |||
| require.True(t, ShouldDisableChannel(constant.ChannelTypeGemini, forbidden)) | |||
| require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, forbidden)) | |||
| invalidKey := types.WithOpenAIError(types.OpenAIError{ | |||
| Message: "bad key", | |||
| Type: "invalid_request_error", | |||
| Code: "invalid_api_key", | |||
| }, http.StatusBadRequest) | |||
| require.True(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, invalidKey)) | |||
| } | |||
| func TestShouldEnableChannelRequiresAutoDisabledStatusWithoutError(t *testing.T) { | |||
| original := common.AutomaticEnableChannelEnabled | |||
| t.Cleanup(func() { common.AutomaticEnableChannelEnabled = original }) | |||
| common.AutomaticEnableChannelEnabled = false | |||
| require.False(t, ShouldEnableChannel(nil, common.ChannelStatusAutoDisabled)) | |||
| common.AutomaticEnableChannelEnabled = true | |||
| require.False(t, ShouldEnableChannel(types.NewError(errors.New("still failing"), types.ErrorCodeBadResponse), common.ChannelStatusAutoDisabled)) | |||
| require.False(t, ShouldEnableChannel(nil, common.ChannelStatusManuallyDisabled)) | |||
| require.True(t, ShouldEnableChannel(nil, common.ChannelStatusAutoDisabled)) | |||
| } | |||
| @@ -120,12 +120,12 @@ func refreshCodexOAuthToken( | |||
| ExpiresIn int `json:"expires_in"` | |||
| } | |||
| if err := common.DecodeJson(resp.Body, &payload); err != nil { | |||
| return nil, err | |||
| } | |||
| if resp.StatusCode < 200 || resp.StatusCode >= 300 { | |||
| return nil, fmt.Errorf("codex oauth refresh failed: status=%d", resp.StatusCode) | |||
| } | |||
| if err := common.DecodeJson(resp.Body, &payload); err != nil { | |||
| return nil, err | |||
| } | |||
| if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 { | |||
| return nil, errors.New("codex oauth refresh response missing fields") | |||
| @@ -181,12 +181,12 @@ func exchangeCodexAuthorizationCode( | |||
| RefreshToken string `json:"refresh_token"` | |||
| ExpiresIn int `json:"expires_in"` | |||
| } | |||
| if err := common.DecodeJson(resp.Body, &payload); err != nil { | |||
| return nil, err | |||
| } | |||
| if resp.StatusCode < 200 || resp.StatusCode >= 300 { | |||
| return nil, fmt.Errorf("codex oauth code exchange failed: status=%d", resp.StatusCode) | |||
| } | |||
| if err := common.DecodeJson(resp.Body, &payload); err != nil { | |||
| return nil, err | |||
| } | |||
| if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 { | |||
| return nil, errors.New("codex oauth token response missing fields") | |||
| } | |||
| @@ -0,0 +1,192 @@ | |||
| package service | |||
| import ( | |||
| "context" | |||
| "encoding/base64" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "net/url" | |||
| "strings" | |||
| "testing" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| ) | |||
| func TestRefreshCodexOAuthTokenReportsHTTPStatusForNonJSONError(t *testing.T) { | |||
| t.Parallel() | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| w.WriteHeader(http.StatusUnauthorized) | |||
| _, _ = w.Write([]byte("upstream unavailable")) | |||
| })) | |||
| defer server.Close() | |||
| _, err := refreshCodexOAuthToken(context.Background(), server.Client(), server.URL, "client-id", "refresh-token") | |||
| if err == nil { | |||
| t.Fatal("expected refresh failure") | |||
| } | |||
| if !strings.Contains(err.Error(), "status=401") { | |||
| t.Fatalf("expected error to contain upstream status, got %q", err) | |||
| } | |||
| } | |||
| func TestRefreshCodexOAuthTokenSendsRefreshGrantAndParsesResponse(t *testing.T) { | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| if r.Method != http.MethodPost || r.Header.Get("Content-Type") != "application/x-www-form-urlencoded" { | |||
| t.Fatalf("unexpected request: method=%s content-type=%q", r.Method, r.Header.Get("Content-Type")) | |||
| } | |||
| if err := r.ParseForm(); err != nil { | |||
| t.Fatalf("parse form: %v", err) | |||
| } | |||
| if r.Form.Get("grant_type") != "refresh_token" || r.Form.Get("client_id") != "client-id" || r.Form.Get("refresh_token") != "refresh-token" { | |||
| t.Fatalf("unexpected refresh form: %#v", r.Form) | |||
| } | |||
| _, _ = w.Write([]byte(`{"access_token":"access","refresh_token":"next-refresh","expires_in":60}`)) | |||
| })) | |||
| defer server.Close() | |||
| before := time.Now() | |||
| result, err := refreshCodexOAuthToken(context.Background(), server.Client(), server.URL, "client-id", " refresh-token ") | |||
| if err != nil { | |||
| t.Fatalf("refresh token: %v", err) | |||
| } | |||
| if result.AccessToken != "access" || result.RefreshToken != "next-refresh" || !result.ExpiresAt.After(before.Add(59*time.Second)) { | |||
| t.Fatalf("unexpected refresh result: %#v", result) | |||
| } | |||
| } | |||
| func TestExchangeCodexAuthorizationCodeSendsPKCEForm(t *testing.T) { | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| if err := r.ParseForm(); err != nil { | |||
| t.Fatalf("parse form: %v", err) | |||
| } | |||
| if r.Form.Get("grant_type") != "authorization_code" || r.Form.Get("code") != "code" || r.Form.Get("code_verifier") != "verifier" || r.Form.Get("redirect_uri") != "http://localhost/callback" { | |||
| t.Fatalf("unexpected authorization-code form: %#v", r.Form) | |||
| } | |||
| _, _ = w.Write([]byte(`{"access_token":"access","refresh_token":"refresh","expires_in":60}`)) | |||
| })) | |||
| defer server.Close() | |||
| result, err := exchangeCodexAuthorizationCode(context.Background(), server.Client(), server.URL, "client-id", " code ", " verifier ", "http://localhost/callback") | |||
| if err != nil { | |||
| t.Fatalf("exchange authorization code: %v", err) | |||
| } | |||
| if result.AccessToken != "access" || result.RefreshToken != "refresh" { | |||
| t.Fatalf("unexpected exchange result: %#v", result) | |||
| } | |||
| } | |||
| func TestExtractCodexClaimsFromJWT(t *testing.T) { | |||
| claims := map[string]any{ | |||
| "email": " user@example.com ", | |||
| codexJWTClaimPath: map[string]any{ | |||
| "chatgpt_account_id": " account-123 ", | |||
| }, | |||
| } | |||
| token := newTestJWT(t, claims) | |||
| accountID, ok := ExtractCodexAccountIDFromJWT(token) | |||
| if !ok || accountID != "account-123" { | |||
| t.Fatalf("expected trimmed account ID, got %q, %t", accountID, ok) | |||
| } | |||
| email, ok := ExtractEmailFromJWT(token) | |||
| if !ok || email != "user@example.com" { | |||
| t.Fatalf("expected trimmed email, got %q, %t", email, ok) | |||
| } | |||
| } | |||
| func TestExtractCodexClaimsRejectMalformedOrEmptyValues(t *testing.T) { | |||
| if _, ok := ExtractCodexAccountIDFromJWT("not-a-jwt"); ok { | |||
| t.Fatal("malformed token must not yield account ID") | |||
| } | |||
| if _, ok := ExtractEmailFromJWT(newTestJWT(t, map[string]any{"email": " "})); ok { | |||
| t.Fatal("empty email must not be accepted") | |||
| } | |||
| if _, ok := ExtractCodexAccountIDFromJWT(newTestJWT(t, map[string]any{ | |||
| codexJWTClaimPath: map[string]any{"chatgpt_account_id": ""}, | |||
| })); ok { | |||
| t.Fatal("empty account ID must not be accepted") | |||
| } | |||
| } | |||
| func TestCreateCodexOAuthAuthorizationFlowUsesPKCEAndState(t *testing.T) { | |||
| flow, err := CreateCodexOAuthAuthorizationFlow() | |||
| if err != nil { | |||
| t.Fatalf("create authorization flow: %v", err) | |||
| } | |||
| if len(flow.State) != 32 { | |||
| t.Fatalf("expected 32-character state, got %q", flow.State) | |||
| } | |||
| if flow.Verifier == "" || flow.Challenge == "" { | |||
| t.Fatal("expected PKCE verifier and challenge") | |||
| } | |||
| u, err := url.Parse(flow.AuthorizeURL) | |||
| if err != nil { | |||
| t.Fatalf("parse authorize URL: %v", err) | |||
| } | |||
| q := u.Query() | |||
| if q.Get("state") != flow.State || q.Get("code_challenge") != flow.Challenge { | |||
| t.Fatalf("authorize URL does not include generated state and challenge: %s", flow.AuthorizeURL) | |||
| } | |||
| if q.Get("code_challenge_method") != "S256" || q.Get("redirect_uri") != codexOAuthRedirectURI { | |||
| t.Fatalf("unexpected PKCE or redirect parameters: %s", flow.AuthorizeURL) | |||
| } | |||
| } | |||
| func TestFetchCodexWhamUsageSendsRequiredHeaders(t *testing.T) { | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| if r.URL.Path != "/backend-api/wham/usage" { | |||
| t.Fatalf("unexpected path: %s", r.URL.Path) | |||
| } | |||
| if r.Header.Get("Authorization") != "Bearer access-token" { | |||
| t.Fatalf("unexpected authorization: %q", r.Header.Get("Authorization")) | |||
| } | |||
| if r.Header.Get("chatgpt-account-id") != "account-123" || r.Header.Get("originator") != "codex_cli_rs" { | |||
| t.Fatalf("missing Codex headers: %#v", r.Header) | |||
| } | |||
| w.Header().Set("Content-Type", "application/json") | |||
| _, _ = w.Write([]byte(`{"limit": 1}`)) | |||
| })) | |||
| defer server.Close() | |||
| status, body, err := FetchCodexWhamUsage(context.Background(), server.Client(), server.URL+"/", " access-token ", " account-123 ") | |||
| if err != nil { | |||
| t.Fatalf("fetch usage: %v", err) | |||
| } | |||
| if status != http.StatusOK || string(body) != `{"limit": 1}` { | |||
| t.Fatalf("unexpected usage response: status=%d body=%q", status, body) | |||
| } | |||
| } | |||
| func TestFetchCodexWhamUsageRejectsMissingInputs(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| client *http.Client | |||
| baseURL string | |||
| accessKey string | |||
| accountID string | |||
| }{ | |||
| {name: "nil client", baseURL: "https://example.com", accessKey: "token", accountID: "account"}, | |||
| {name: "empty base URL", client: http.DefaultClient, accessKey: "token", accountID: "account"}, | |||
| {name: "empty access token", client: http.DefaultClient, baseURL: "https://example.com", accountID: "account"}, | |||
| {name: "empty account ID", client: http.DefaultClient, baseURL: "https://example.com", accessKey: "token"}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| if _, _, err := FetchCodexWhamUsage(context.Background(), tt.client, tt.baseURL, tt.accessKey, tt.accountID); err == nil { | |||
| t.Fatal("expected input validation error") | |||
| } | |||
| }) | |||
| } | |||
| } | |||
| func newTestJWT(t *testing.T, claims map[string]any) string { | |||
| t.Helper() | |||
| payload, err := common.Marshal(claims) | |||
| if err != nil { | |||
| t.Fatalf("marshal claims: %v", err) | |||
| } | |||
| return "header." + base64.RawURLEncoding.EncodeToString(payload) + ".signature" | |||
| } | |||
| @@ -0,0 +1,31 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestConvertGeminiRoleToOpenAIUsesSafeUserFallback(t *testing.T) { | |||
| require.Equal(t, "user", convertGeminiRoleToOpenAI("user")) | |||
| require.Equal(t, "assistant", convertGeminiRoleToOpenAI("model")) | |||
| require.Equal(t, "function", convertGeminiRoleToOpenAI("function")) | |||
| require.Equal(t, "user", convertGeminiRoleToOpenAI("unknown")) | |||
| } | |||
| func TestExtractTextFromGeminiPartsSkipsEmptyParts(t *testing.T) { | |||
| parts := []dto.GeminiPart{{Text: "first"}, {}, {Text: "last"}} | |||
| require.Equal(t, "first\nlast", extractTextFromGeminiParts(parts)) | |||
| } | |||
| func TestToJSONStringSerializesValuesAndUsesObjectFallback(t *testing.T) { | |||
| require.Equal(t, `{"answer":42}`, toJSONString(map[string]int{"answer": 42})) | |||
| require.Equal(t, "{}", toJSONString(make(chan int))) | |||
| } | |||
| func TestStopReasonOpenAI2Claude(t *testing.T) { | |||
| require.Equal(t, "end_turn", stopReasonOpenAI2Claude("stop")) | |||
| require.NotEmpty(t, stopReasonOpenAI2Claude("unknown")) | |||
| } | |||
| @@ -42,7 +42,7 @@ func setupDoubaoAssetChannelDB(t *testing.T) *gorm.DB { | |||
| sqlDB, err := db.DB() | |||
| require.NoError(t, err) | |||
| sqlDB.SetMaxOpenConns(1) | |||
| require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{})) | |||
| require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.Token{})) | |||
| t.Cleanup(func() { | |||
| _ = sqlDB.Close() | |||
| @@ -0,0 +1,47 @@ | |||
| package service | |||
| import ( | |||
| "net/http" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestSmartDetectMimeTypeUsesTrustedHeaderBeforeFilenameAndURL(t *testing.T) { | |||
| resp := &http.Response{Header: http.Header{ | |||
| "Content-Type": []string{"image/png; charset=binary"}, | |||
| "Content-Disposition": []string{`attachment; filename="report.pdf"`}, | |||
| }} | |||
| require.Equal(t, "image/png", smartDetectMimeType(resp, "https://example.com/video.mp4", []byte("not-an-image"))) | |||
| } | |||
| func TestSmartDetectMimeTypeFallsBackToFilenameAndURL(t *testing.T) { | |||
| filenameResp := &http.Response{Header: http.Header{ | |||
| "Content-Type": []string{"application/octet-stream"}, | |||
| "Content-Disposition": []string{`attachment; filename="photo.JPEG"`}, | |||
| }} | |||
| urlResp := &http.Response{Header: http.Header{"Content-Type": []string{"application/octet-stream"}}} | |||
| require.Equal(t, "image/jpeg", smartDetectMimeType(filenameResp, "https://example.com/archive.bin", nil)) | |||
| require.Equal(t, "application/pdf", smartDetectMimeType(urlResp, "https://example.com/report.PDF?download=1", nil)) | |||
| } | |||
| func TestLoadFromBase64HonorsExplicitMimeTypeAndRejectsInvalidData(t *testing.T) { | |||
| cached, err := loadFromBase64("data:text/plain;base64,aGVsbG8=", "application/custom") | |||
| require.NoError(t, err) | |||
| require.Equal(t, "application/custom", cached.MimeType) | |||
| require.EqualValues(t, 5, cached.Size) | |||
| cached.Close() | |||
| _, err = loadFromBase64("not-base64", "") | |||
| require.Error(t, err) | |||
| } | |||
| func TestDetectFileType(t *testing.T) { | |||
| require.Equal(t, types.FileTypeImage, DetectFileType("image/webp")) | |||
| require.Equal(t, types.FileTypeAudio, DetectFileType("audio/mpeg")) | |||
| require.Equal(t, types.FileTypeVideo, DetectFileType("video/mp4")) | |||
| require.Equal(t, types.FileTypeFile, DetectFileType("application/pdf")) | |||
| } | |||
| @@ -0,0 +1,41 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetUserUsableGroupsAppliesSpecialRulesAndKeepsOwnGroup(t *testing.T) { | |||
| originalGroups := setting.UserUsableGroups2JSONString() | |||
| originalSpecial := ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup.MarshalJSONString() | |||
| t.Cleanup(func() { | |||
| require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalGroups)) | |||
| require.NoError(t, types.LoadFromJsonString(ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup, originalSpecial)) | |||
| }) | |||
| require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","remove":"Remove"}`)) | |||
| require.NoError(t, types.LoadFromJsonString(ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup, `{"staff":{"-:remove":"","+:premium":"Premium"}}`)) | |||
| groups := GetUserUsableGroups("staff") | |||
| require.Equal(t, "Default", groups["default"]) | |||
| require.NotContains(t, groups, "remove") | |||
| require.Equal(t, "Premium", groups["premium"]) | |||
| require.Equal(t, "用户分组", groups["staff"]) | |||
| } | |||
| func TestGetUserAutoGroupFiltersConfiguredGroupsByAccess(t *testing.T) { | |||
| originalGroups := setting.UserUsableGroups2JSONString() | |||
| originalAutoGroups := setting.AutoGroups2JsonString() | |||
| t.Cleanup(func() { | |||
| require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalGroups)) | |||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) | |||
| }) | |||
| require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) | |||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["vip","missing","default"]`)) | |||
| require.Equal(t, []string{"vip", "default"}, GetUserAutoGroup("default")) | |||
| } | |||
| @@ -0,0 +1,64 @@ | |||
| package service | |||
| import ( | |||
| "net/http" | |||
| "testing" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func preserveHTTPClientState(t *testing.T) { | |||
| t.Helper() | |||
| oldDefaultClient := httpClient | |||
| oldTimeout := common.RelayTimeout | |||
| proxyClientLock.Lock() | |||
| oldProxyClients := proxyClients | |||
| proxyClients = make(map[string]*http.Client) | |||
| proxyClientLock.Unlock() | |||
| t.Cleanup(func() { | |||
| httpClient = oldDefaultClient | |||
| common.RelayTimeout = oldTimeout | |||
| proxyClientLock.Lock() | |||
| proxyClients = oldProxyClients | |||
| proxyClientLock.Unlock() | |||
| }) | |||
| } | |||
| func TestNewProxyHttpClientCachesHTTPProxyAndRejectsUnsupportedScheme(t *testing.T) { | |||
| preserveHTTPClientState(t) | |||
| common.RelayTimeout = 7 | |||
| first, err := NewProxyHttpClient("http://proxy.example:8080") | |||
| require.NoError(t, err) | |||
| second, err := NewProxyHttpClient("http://proxy.example:8080") | |||
| require.NoError(t, err) | |||
| require.Same(t, first, second) | |||
| require.Equal(t, 7*time.Second, first.Timeout) | |||
| _, err = NewProxyHttpClient("ftp://proxy.example:21") | |||
| require.ErrorContains(t, err, "unsupported proxy scheme") | |||
| } | |||
| func TestNewProxyHttpClientReturnsInitializedDefaultClientForEmptyProxy(t *testing.T) { | |||
| preserveHTTPClientState(t) | |||
| defaultClient := &http.Client{} | |||
| httpClient = defaultClient | |||
| client, err := NewProxyHttpClient("") | |||
| require.NoError(t, err) | |||
| require.Same(t, defaultClient, client) | |||
| } | |||
| func TestInitHttpClientAppliesConfiguredTimeout(t *testing.T) { | |||
| preserveHTTPClientState(t) | |||
| common.RelayTimeout = 3 | |||
| InitHttpClient() | |||
| require.NotNil(t, GetHttpClient()) | |||
| require.Equal(t, 3*time.Second, GetHttpClient().Timeout) | |||
| require.NotNil(t, GetHttpClient().CheckRedirect) | |||
| } | |||
| @@ -0,0 +1,45 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestDecodeBase64FileDataPreservesDeclaredMimeType(t *testing.T) { | |||
| mimeType, payload, err := DecodeBase64FileData("data:application/pdf;base64,SGVsbG8=") | |||
| require.NoError(t, err) | |||
| require.Equal(t, "application/pdf", mimeType) | |||
| require.Equal(t, "SGVsbG8=", payload) | |||
| } | |||
| func TestDecodeBase64ImageDataRejectsEmptyAndMalformedData(t *testing.T) { | |||
| _, _, _, err := DecodeBase64ImageData("") | |||
| require.Error(t, err) | |||
| _, _, _, err = DecodeBase64ImageData("data:image/png;base64,not-base64") | |||
| require.Error(t, err) | |||
| } | |||
| func TestDecodeBase64ImageDataParsesPNGDataURL(t *testing.T) { | |||
| const pngBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVQIHWP4z8DwHwAFgAI/ScLxZQAAAABJRU5ErkJggg==" | |||
| config, format, payload, err := DecodeBase64ImageData("data:image/png;base64," + pngBase64) | |||
| require.NoError(t, err) | |||
| require.Equal(t, 1, config.Width) | |||
| require.Equal(t, 1, config.Height) | |||
| require.Equal(t, "png", format) | |||
| require.Equal(t, pngBase64, payload) | |||
| mimeType, inferredPayload, err := DecodeBase64FileData(pngBase64) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "image/png", mimeType) | |||
| require.Equal(t, pngBase64, inferredPayload) | |||
| } | |||
| func TestDecodeBase64FileDataFallsBackToImageDetectionForBareData(t *testing.T) { | |||
| _, _, err := DecodeBase64FileData("not-base64") | |||
| require.Error(t, err) | |||
| } | |||
| @@ -140,6 +140,9 @@ func ConvertSimpleChangeParams(content string) *dto.MidjourneyRequest { | |||
| } | |||
| action := strings.ToLower(split[1]) | |||
| if action == "" { | |||
| return nil | |||
| } | |||
| changeParams := &dto.MidjourneyRequest{} | |||
| changeParams.TaskId = split[0] | |||
| @@ -0,0 +1,25 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestConvertSimpleChangeParamsParsesSupportedActions(t *testing.T) { | |||
| upscale := ConvertSimpleChangeParams("task-1 U3") | |||
| require.NotNil(t, upscale) | |||
| require.Equal(t, "task-1", upscale.TaskId) | |||
| require.Equal(t, "UPSCALE", upscale.Action) | |||
| require.Equal(t, 3, upscale.Index) | |||
| reroll := ConvertSimpleChangeParams("task-1 R") | |||
| require.NotNil(t, reroll) | |||
| require.Equal(t, "REROLL", reroll.Action) | |||
| } | |||
| func TestConvertSimpleChangeParamsRejectsBlankOrMalformedAction(t *testing.T) { | |||
| require.Nil(t, ConvertSimpleChangeParams("task-1 ")) | |||
| require.Nil(t, ConvertSimpleChangeParams("task-1 X1")) | |||
| require.Nil(t, ConvertSimpleChangeParams("task-1 U0")) | |||
| } | |||
| @@ -0,0 +1,54 @@ | |||
| package openaicompat | |||
| import ( | |||
| "github.com/QuantumNous/new-api/common" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/setting/model_setting" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestShouldChatCompletionsUseResponsesPolicyRequiresEnabledChannelAndMatchingModel(t *testing.T) { | |||
| policy := model_setting.ChatCompletionsToResponsesPolicy{ | |||
| Enabled: true, | |||
| ChannelIDs: []int{42}, | |||
| ModelPatterns: []string{`^gpt-5(?:-mini)?$`, `[`}, | |||
| } | |||
| require.True(t, ShouldChatCompletionsUseResponsesPolicy(policy, 42, 0, "gpt-5-mini")) | |||
| require.False(t, ShouldChatCompletionsUseResponsesPolicy(policy, 42, 0, "gpt-4o")) | |||
| require.False(t, ShouldChatCompletionsUseResponsesPolicy(policy, 7, 0, "gpt-5")) | |||
| } | |||
| func TestNormalizeChatImageURLToStringSupportsSDKAndDecodedForms(t *testing.T) { | |||
| require.Equal(t, "https://example.com/image.png", normalizeChatImageURLToString("https://example.com/image.png")) | |||
| require.Equal(t, "https://example.com/image.png", normalizeChatImageURLToString(map[string]any{"url": "https://example.com/image.png"})) | |||
| require.Equal(t, "https://example.com/image.png", normalizeChatImageURLToString(&dto.MessageImageUrl{Url: "https://example.com/image.png"})) | |||
| unknown := map[string]any{"other": "value"} | |||
| require.Equal(t, unknown, normalizeChatImageURLToString(unknown)) | |||
| } | |||
| func TestConvertChatResponseFormatToResponsesTextFlattensJSONSchema(t *testing.T) { | |||
| input := &dto.ResponseFormat{Type: "json_schema", JsonSchema: []byte(`{"name":"answer","json_schema":{"schema":{"type":"object"},"strict":true}}`)} | |||
| raw := convertChatResponseFormatToResponsesText(input) | |||
| var payload map[string]map[string]any | |||
| require.NoError(t, common.Unmarshal(raw, &payload)) | |||
| format := payload["format"] | |||
| require.Equal(t, "json_schema", format["type"]) | |||
| require.Equal(t, "answer", format["name"]) | |||
| require.Equal(t, true, format["strict"]) | |||
| require.Equal(t, map[string]any{"type": "object"}, format["schema"]) | |||
| } | |||
| func TestShouldChatCompletionsUseResponsesPolicyAllowsConfiguredChannelType(t *testing.T) { | |||
| policy := model_setting.ChatCompletionsToResponsesPolicy{ | |||
| Enabled: true, | |||
| ChannelTypes: []int{9}, | |||
| ModelPatterns: []string{`^o[0-9]+$`}, | |||
| } | |||
| require.True(t, ShouldChatCompletionsUseResponsesPolicy(policy, 0, 9, "o3")) | |||
| } | |||
| @@ -0,0 +1,54 @@ | |||
| package passkey | |||
| import ( | |||
| "crypto/tls" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestResolveOriginsRejectsInsecureConfiguredOriginUnlessAllowed(t *testing.T) { | |||
| request := httptest.NewRequest("GET", "https://gateway.example", nil) | |||
| settings := &system_setting.PasskeySettings{Origins: "http://gateway.example", AllowInsecureOrigin: false} | |||
| _, err := resolveOrigins(request, settings) | |||
| require.Error(t, err) | |||
| settings.AllowInsecureOrigin = true | |||
| origins, err := resolveOrigins(request, settings) | |||
| require.NoError(t, err) | |||
| require.Equal(t, []string{"http://gateway.example"}, origins) | |||
| } | |||
| func TestResolveOriginsUsesForwardedHTTPSAndRejectsUnsecuredPublicHost(t *testing.T) { | |||
| settings := &system_setting.PasskeySettings{} | |||
| request := httptest.NewRequest("GET", "http://gateway.example:8443", nil) | |||
| request.Header.Set("X-Forwarded-Proto", "https, http") | |||
| origins, err := resolveOrigins(request, settings) | |||
| require.NoError(t, err) | |||
| require.Equal(t, []string{"https://gateway.example:8443"}, origins) | |||
| request = httptest.NewRequest("GET", "http://gateway.example", nil) | |||
| _, err = resolveOrigins(request, settings) | |||
| require.Error(t, err) | |||
| } | |||
| func TestResolveRPIDAndSchemeDetection(t *testing.T) { | |||
| settings := &system_setting.PasskeySettings{} | |||
| rpID, err := resolveRPID(nil, settings, []string{"https://gateway.example:8443"}) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "gateway.example", rpID) | |||
| settings.RPID = " configured.example:9443 " | |||
| rpID, err = resolveRPID(nil, settings, nil) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "configured.example", rpID) | |||
| request := httptest.NewRequest("GET", "http://localhost", nil) | |||
| request.TLS = &tls.ConnectionState{} | |||
| require.Equal(t, "https", detectScheme(request)) | |||
| require.Equal(t, "::1", hostWithoutPort("[::1]:443")) | |||
| } | |||
| @@ -0,0 +1,20 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/setting" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestSensitiveWordReplaceMasksChineseWordAfterASCII(t *testing.T) { | |||
| oldWords := setting.SensitiveWords | |||
| setting.SensitiveWords = []string{"秘密"} | |||
| t.Cleanup(func() { setting.SensitiveWords = oldWords }) | |||
| found, words, replaced := SensitiveWordReplace("prefix秘密suffix", false) | |||
| require.True(t, found) | |||
| require.Equal(t, []string{"秘密"}, words) | |||
| require.Equal(t, "prefix**###**suffix", replaced) | |||
| } | |||
| @@ -12,20 +12,29 @@ import ( | |||
| ) | |||
| func SundaySearch(text string, pattern string) bool { | |||
| textRunes := []rune(text) | |||
| patternRunes := []rune(pattern) | |||
| if len(patternRunes) == 0 { | |||
| return true | |||
| } | |||
| if len(patternRunes) > len(textRunes) { | |||
| return false | |||
| } | |||
| // 计算偏移表 | |||
| offset := make(map[rune]int) | |||
| for i, c := range pattern { | |||
| offset[c] = len(pattern) - i | |||
| for i, c := range patternRunes { | |||
| offset[c] = len(patternRunes) - i | |||
| } | |||
| // 文本串长度和模式串长度 | |||
| n, m := len(text), len(pattern) | |||
| n, m := len(textRunes), len(patternRunes) | |||
| // 主循环,i表示当前对齐的文本串位置 | |||
| for i := 0; i <= n-m; { | |||
| // 检查子串 | |||
| j := 0 | |||
| for j < m && text[i+j] == pattern[j] { | |||
| for j < m && textRunes[i+j] == patternRunes[j] { | |||
| j++ | |||
| } | |||
| // 如果完全匹配,返回匹配位置 | |||
| @@ -35,7 +44,7 @@ func SundaySearch(text string, pattern string) bool { | |||
| // 如果还有剩余字符,则检查下一位字符在偏移表中的值 | |||
| if i+m < n { | |||
| next := rune(text[i+m]) | |||
| next := textRunes[i+m] | |||
| if val, ok := offset[next]; ok { | |||
| i += val // 存在于偏移表中,进行跳跃 | |||
| } else { | |||
| @@ -0,0 +1,20 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestSundaySearchFindsUnicodePatternAfterUnmatchedPrefix(t *testing.T) { | |||
| require.True(t, SundaySearch("你好世界", "世界")) | |||
| } | |||
| func TestSundaySearchHandlesASCIIMatchesAndMisses(t *testing.T) { | |||
| require.True(t, SundaySearch("prefix-target-suffix", "target")) | |||
| require.False(t, SundaySearch("prefix-target-suffix", "missing")) | |||
| } | |||
| func TestRemoveDuplicatePreservesFirstOccurrenceOrder(t *testing.T) { | |||
| require.Equal(t, []string{"a", "b", "c"}, RemoveDuplicate([]string{"a", "b", "a", "c", "b"})) | |||
| } | |||
| @@ -0,0 +1,35 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestEstimateTokenClassifiesWordNumberAndUnicodeBoundaries(t *testing.T) { | |||
| tests := []struct { | |||
| name string | |||
| text string | |||
| want int | |||
| }{ | |||
| {name: "empty", text: "", want: 0}, | |||
| {name: "latin word", text: "hello", want: 2}, | |||
| {name: "latin number transition", text: "v2", want: 3}, | |||
| {name: "cjk", text: "你", want: 1}, | |||
| {name: "emoji", text: "🙂", want: 3}, | |||
| {name: "url delimiter", text: "/", want: 1}, | |||
| } | |||
| for _, tt := range tests { | |||
| t.Run(tt.name, func(t *testing.T) { | |||
| require.Equal(t, tt.want, EstimateToken(OpenAI, tt.text)) | |||
| }) | |||
| } | |||
| } | |||
| func TestEstimateTokenByModelSelectsProviderAndEmptyInputIsFree(t *testing.T) { | |||
| require.Equal(t, 0, EstimateTokenByModel("gpt-5", "")) | |||
| require.Equal(t, EstimateToken(Gemini, "hello"), EstimateTokenByModel("GEMINI-2.5-PRO", "hello")) | |||
| require.Equal(t, EstimateToken(Claude, "hello"), EstimateTokenByModel("claude-sonnet", "hello")) | |||
| require.Equal(t, EstimateToken(OpenAI, "hello"), EstimateTokenByModel("other-model", "hello")) | |||
| } | |||
| @@ -0,0 +1,30 @@ | |||
| package service | |||
| import ( | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/gin-gonic/gin" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestResponseText2UsageSetsLocalCountFlagAndTotals(t *testing.T) { | |||
| ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) | |||
| usage := ResponseText2Usage(ctx, "hello", "gpt-5", 3) | |||
| require.Equal(t, 3, usage.PromptTokens) | |||
| require.Equal(t, EstimateTokenByModel("gpt-5", "hello"), usage.CompletionTokens) | |||
| require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens) | |||
| require.True(t, common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens)) | |||
| } | |||
| func TestValidUsageRequiresAtLeastOneNonZeroTokenCount(t *testing.T) { | |||
| require.False(t, ValidUsage(nil)) | |||
| require.False(t, ValidUsage(&dto.Usage{})) | |||
| require.True(t, ValidUsage(&dto.Usage{PromptTokens: 1})) | |||
| require.True(t, ValidUsage(&dto.Usage{CompletionTokens: 1})) | |||
| } | |||
| @@ -2,6 +2,7 @@ package service | |||
| import ( | |||
| "errors" | |||
| "sort" | |||
| "strings" | |||
| "github.com/QuantumNous/new-api/common" | |||
| @@ -17,6 +18,10 @@ const ( | |||
| VideoAssetFamilyKling VideoAssetFamily = "kling" | |||
| ) | |||
| func VideoAssetFamilies() []VideoAssetFamily { | |||
| return []VideoAssetFamily{VideoAssetFamilySeedance, VideoAssetFamilyKling} | |||
| } | |||
| func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int { | |||
| switch family { | |||
| case VideoAssetFamilySeedance: | |||
| @@ -107,6 +112,55 @@ func BindVideoAssetChannel(userId int, tokenGroup string, channel *model.Channel | |||
| }) | |||
| } | |||
| func ClearVideoAssetChannelBinding(userId int, tokenGroup string, family VideoAssetFamily) error { | |||
| channelTypes := VideoAssetChannelTypesForFamily(family) | |||
| if len(channelTypes) == 0 { | |||
| return nil | |||
| } | |||
| return model.DB.Transaction(func(tx *gorm.DB) error { | |||
| return model.DeleteUserAssetChannelsByTypesWithTx(tx, userId, channelTypes, tokenGroup) | |||
| }) | |||
| } | |||
| func GetVideoAssetChannelCandidates(tokenGroup string, family VideoAssetFamily) ([]*model.Channel, error) { | |||
| channelTypes := VideoAssetChannelTypesForFamily(family) | |||
| if len(channelTypes) == 0 { | |||
| return nil, nil | |||
| } | |||
| candidates := make(map[int]*model.Channel) | |||
| for _, channelType := range channelTypes { | |||
| for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { | |||
| channels, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| for _, candidate := range channels { | |||
| if candidate == nil || candidate.Status != common.ChannelStatusEnabled || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { | |||
| continue | |||
| } | |||
| channel, err := model.GetChannelById(candidate.Id, true) | |||
| if err != nil { | |||
| return nil, err | |||
| } | |||
| if channel.Status == common.ChannelStatusEnabled && strings.TrimSpace(channel.Key) != "" && MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) { | |||
| channel.Key = "" | |||
| channel.Keys = nil | |||
| candidates[channel.Id] = channel | |||
| } | |||
| } | |||
| if len(channels) < DoubaoAssetChannelPageSize { | |||
| break | |||
| } | |||
| } | |||
| } | |||
| result := make([]*model.Channel, 0, len(candidates)) | |||
| for _, channel := range candidates { | |||
| result = append(result, channel) | |||
| } | |||
| sort.Slice(result, func(i, j int) bool { return result[i].Id < result[j].Id }) | |||
| return result, nil | |||
| } | |||
| func HasVideoAssetChannelForGroupModel(tokenGroup string, modelName string, family VideoAssetFamily) bool { | |||
| tokenGroup = strings.TrimSpace(tokenGroup) | |||
| modelName = strings.TrimSpace(modelName) | |||
| @@ -0,0 +1,55 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/constant" | |||
| "github.com/QuantumNous/new-api/model" | |||
| "github.com/stretchr/testify/assert" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGetUserConcreteTokenGroups(t *testing.T) { | |||
| db := setupDoubaoAssetChannelDB(t) | |||
| for i, group := range []string{"vip", "default", "default", "auto", ""} { | |||
| require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "key-" + string(rune('0'+i)), Group: group}).Error) | |||
| } | |||
| groups, err := model.GetUserConcreteTokenGroups(10) | |||
| require.NoError(t, err) | |||
| assert.Equal(t, []string{"default", "vip"}, groups) | |||
| } | |||
| func TestGetVideoAssetChannelCandidatesFiltersGroupAndStatus(t *testing.T) { | |||
| db := setupDoubaoAssetChannelDB(t) | |||
| createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "key", common.ChannelStatusEnabled) | |||
| createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled) | |||
| createDoubaoAssetChannelForTest(t, db, 17, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusAutoDisabled) | |||
| createDoubaoAssetChannelForTest(t, db, 18, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "vip", "key", common.ChannelStatusEnabled) | |||
| createDoubaoAssetChannelForTest(t, db, 19, constant.ChannelTypeChinaMobileSeedance, "default", "\n", common.ChannelStatusEnabled) | |||
| candidates, err := GetVideoAssetChannelCandidates("default", VideoAssetFamilySeedance) | |||
| require.NoError(t, err) | |||
| require.Len(t, candidates, 2) | |||
| assert.Equal(t, []int{7, 16}, []int{candidates[0].Id, candidates[1].Id}) | |||
| } | |||
| func TestClearVideoAssetChannelBindingClearsOnlyFamily(t *testing.T) { | |||
| db := setupDoubaoAssetChannelDB(t) | |||
| createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "key", common.ChannelStatusEnabled) | |||
| createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) | |||
| require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59)) | |||
| require.NoError(t, ClearVideoAssetChannelBinding(10, "default", VideoAssetFamilySeedance)) | |||
| seedanceBindings, err := model.GetUserAssetChannelsByTypes(10, VideoAssetChannelTypesForFamily(VideoAssetFamilySeedance), "default") | |||
| require.NoError(t, err) | |||
| assert.Empty(t, seedanceBindings) | |||
| klingBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default") | |||
| require.NoError(t, err) | |||
| assert.NotNil(t, klingBinding) | |||
| } | |||
| @@ -0,0 +1,36 @@ | |||
| package service | |||
| import ( | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/types" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestIsViolationFeeCodeRecognizesOnlyReservedPrefix(t *testing.T) { | |||
| require.True(t, IsViolationFeeCode(types.ErrorCode("violation_fee.grok_csam"))) | |||
| require.False(t, IsViolationFeeCode(types.ErrorCode("grok_csam"))) | |||
| } | |||
| func TestCalcViolationFeeQuotaRejectsInvalidInputsAndRounds(t *testing.T) { | |||
| oldQuotaPerUnit := common.QuotaPerUnit | |||
| common.QuotaPerUnit = 1000 | |||
| t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) | |||
| require.Equal(t, 0, calcViolationFeeQuota(0, 1)) | |||
| require.Equal(t, 0, calcViolationFeeQuota(1, 0)) | |||
| require.Equal(t, 333, calcViolationFeeQuota(0.333, 1)) | |||
| require.Equal(t, 500, calcViolationFeeQuota(0.25, 2)) | |||
| } | |||
| func TestNormalizeViolationFeeErrorMarksExistingViolationCodeAsNonRetryable(t *testing.T) { | |||
| err := types.WithOpenAIError(types.OpenAIError{ | |||
| Code: string(types.ErrorCodeViolationFeeGrokCSAM), | |||
| }, 400) | |||
| normalized := NormalizeViolationFeeError(err) | |||
| require.Equal(t, types.ErrorCodeViolationFeeGrokCSAM, normalized.GetErrorCode()) | |||
| require.True(t, types.IsSkipRetryError(normalized)) | |||
| } | |||
| @@ -8,6 +8,7 @@ import ( | |||
| "encoding/json" | |||
| "fmt" | |||
| "net/http" | |||
| "strings" | |||
| "time" | |||
| "github.com/QuantumNous/new-api/common" | |||
| @@ -36,7 +37,7 @@ func SendWebhookNotify(webhookURL string, secret string, data dto.Notify) error | |||
| // 处理占位符 | |||
| content := data.Content | |||
| for _, value := range data.Values { | |||
| content = fmt.Sprintf(content, value) | |||
| content = strings.Replace(content, dto.ContentValueParam, fmt.Sprintf("%v", value), 1) | |||
| } | |||
| // 构建 webhook 负载 | |||
| @@ -0,0 +1,48 @@ | |||
| package service | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "github.com/QuantumNous/new-api/dto" | |||
| "github.com/QuantumNous/new-api/setting/system_setting" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestGenerateSignatureIsStableAndSecretBound(t *testing.T) { | |||
| payload := []byte(`{"event":"quota_exceed"}`) | |||
| require.Equal(t, "e8c1ac446d1fab782a3587e16b8dae9b42afd915f6c56a25702806f35db4ecaf", generateSignature("webhook-secret", payload)) | |||
| require.NotEqual(t, generateSignature("other-secret", payload), generateSignature("webhook-secret", payload)) | |||
| require.NotEqual(t, generateSignature("webhook-secret", []byte(`{"event":"other"}`)), generateSignature("webhook-secret", payload)) | |||
| } | |||
| func TestSendWebhookNotifyReplacesValuePlaceholdersBeforeSigningAndSending(t *testing.T) { | |||
| var received WebhookPayload | |||
| var signature string | |||
| server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| body, err := io.ReadAll(r.Body) | |||
| require.NoError(t, err) | |||
| require.NoError(t, common.Unmarshal(body, &received)) | |||
| signature = r.Header.Get("X-Webhook-Signature") | |||
| w.WriteHeader(http.StatusNoContent) | |||
| })) | |||
| defer server.Close() | |||
| oldClient := httpClient | |||
| httpClient = server.Client() | |||
| t.Cleanup(func() { httpClient = oldClient }) | |||
| setting := system_setting.GetFetchSetting() | |||
| oldSetting := *setting | |||
| *setting = system_setting.FetchSetting{EnableSSRFProtection: false} | |||
| t.Cleanup(func() { *setting = oldSetting }) | |||
| err := SendWebhookNotify(server.URL, "secret", dto.NewNotify("quota", "Quota", "remaining: {{value}}, used: {{value}}", []interface{}{12, 3})) | |||
| require.NoError(t, err) | |||
| require.Equal(t, "remaining: 12, used: 3", received.Content) | |||
| require.NotEmpty(t, signature) | |||
| } | |||
| @@ -20,8 +20,12 @@ func ContainsAutoGroup(group string) bool { | |||
| } | |||
| func UpdateAutoGroupsByJsonString(jsonString string) error { | |||
| autoGroups = make([]string, 0) | |||
| return common.Unmarshal([]byte(jsonString), &autoGroups) | |||
| var groups []string | |||
| if err := common.Unmarshal([]byte(jsonString), &groups); err != nil { | |||
| return err | |||
| } | |||
| autoGroups = groups | |||
| return nil | |||
| } | |||
| func AutoGroups2JsonString() string { | |||
| @@ -31,8 +31,12 @@ var Chats = []map[string]string{ | |||
| } | |||
| func UpdateChatsByJsonString(jsonString string) error { | |||
| Chats = make([]map[string]string, 0) | |||
| return json.Unmarshal([]byte(jsonString), &Chats) | |||
| var chats []map[string]string | |||
| if err := json.Unmarshal([]byte(jsonString), &chats); err != nil { | |||
| return err | |||
| } | |||
| Chats = chats | |||
| return nil | |||
| } | |||
| func Chats2JsonString() string { | |||
| @@ -0,0 +1,29 @@ | |||
| package setting | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestUpdateAutoGroupsByJsonStringKeepsExistingGroupsOnInvalidJSON(t *testing.T) { | |||
| old := AutoGroups2JsonString() | |||
| t.Cleanup(func() { require.NoError(t, UpdateAutoGroupsByJsonString(old)) }) | |||
| require.NoError(t, UpdateAutoGroupsByJsonString(`["default"]`)) | |||
| err := UpdateAutoGroupsByJsonString("not-json") | |||
| require.Error(t, err) | |||
| require.Equal(t, []string{"default"}, GetAutoGroups()) | |||
| } | |||
| func TestUpdateChatsByJsonStringKeepsExistingChatsOnInvalidJSON(t *testing.T) { | |||
| old := Chats2JsonString() | |||
| t.Cleanup(func() { require.NoError(t, UpdateChatsByJsonString(old)) }) | |||
| require.NoError(t, UpdateChatsByJsonString(`[{"name":"url"}]`)) | |||
| err := UpdateChatsByJsonString("not-json") | |||
| require.Error(t, err) | |||
| require.Equal(t, []map[string]string{{"name": "url"}}, Chats) | |||
| } | |||
| @@ -37,8 +37,12 @@ var PayMethods = []map[string]string{ | |||
| } | |||
| func UpdatePayMethodsByJsonString(jsonString string) error { | |||
| PayMethods = make([]map[string]string, 0) | |||
| return common.Unmarshal([]byte(jsonString), &PayMethods) | |||
| var methods []map[string]string | |||
| if err := common.Unmarshal([]byte(jsonString), &methods); err != nil { | |||
| return err | |||
| } | |||
| PayMethods = methods | |||
| return nil | |||
| } | |||
| func PayMethods2JsonString() string { | |||
| @@ -0,0 +1,18 @@ | |||
| package operation_setting | |||
| import ( | |||
| "testing" | |||
| "github.com/stretchr/testify/require" | |||
| ) | |||
| func TestUpdatePayMethodsByJsonStringKeepsExistingMethodsOnInvalidJSON(t *testing.T) { | |||
| old := PayMethods2JsonString() | |||
| t.Cleanup(func() { require.NoError(t, UpdatePayMethodsByJsonString(old)) }) | |||
| require.NoError(t, UpdatePayMethodsByJsonString(`[{"type":"card"}]`)) | |||
| err := UpdatePayMethodsByJsonString("not-json") | |||
| require.Error(t, err) | |||
| require.Equal(t, []map[string]string{{"type": "card"}}, PayMethods) | |||
| } | |||
| @@ -56,6 +56,7 @@ import { | |||
| IconPlus, | |||
| } from '@douyinfe/semi-icons'; | |||
| import UserBindingManagementModal from './UserBindingManagementModal'; | |||
| import UserVideoChannelBindingModal from './UserVideoChannelBindingModal'; | |||
| import UserModelRateLimitSection from './UserModelRateLimitSection'; | |||
| const { Text, Title } = Typography; | |||
| @@ -70,6 +71,7 @@ const EditUserModal = (props) => { | |||
| const isMobile = useIsMobile(); | |||
| const [groupOptions, setGroupOptions] = useState([]); | |||
| const [bindingModalVisible, setBindingModalVisible] = useState(false); | |||
| const [videoBindingModalVisible, setVideoBindingModalVisible] = useState(false); | |||
| const formApiRef = useRef(null); | |||
| const [isSyncedUser, setIsSyncedUser] = useState(false); | |||
| @@ -122,6 +124,7 @@ const EditUserModal = (props) => { | |||
| loadUser(); | |||
| if (userId) fetchGroups(); | |||
| setBindingModalVisible(false); | |||
| setVideoBindingModalVisible(false); | |||
| }, [props.editingUser.id]); | |||
| const openBindingModal = () => { | |||
| @@ -132,6 +135,9 @@ const EditUserModal = (props) => { | |||
| setBindingModalVisible(false); | |||
| }; | |||
| const openVideoBindingModal = () => setVideoBindingModalVisible(true); | |||
| const closeVideoBindingModal = () => setVideoBindingModalVisible(false); | |||
| /* ----------------------- submit ----------------------- */ | |||
| const submit = async (values) => { | |||
| setLoading(true); | |||
| @@ -284,6 +290,26 @@ const EditUserModal = (props) => { | |||
| </Card> | |||
| {/* 权限设置 */} | |||
| {userId && ( | |||
| <Card className='!rounded-2xl shadow-sm border-0'> | |||
| <div className='flex items-center justify-between gap-3'> | |||
| <div className='flex items-center min-w-0'> | |||
| <Avatar size='small' color='orange' className='mr-2 shadow-md'> | |||
| <IconLink size={16} /> | |||
| </Avatar> | |||
| <div className='min-w-0'> | |||
| <Text className='text-lg font-medium'>{t('视频渠道绑定')}</Text> | |||
| <div className='text-xs text-gray-600'> | |||
| {t('按该用户 Token 分组配置 Seedance、Kling 等视频路由')} | |||
| </div> | |||
| </div> | |||
| </div> | |||
| <Button type='primary' theme='outline' onClick={openVideoBindingModal}> | |||
| {t('配置视频渠道')} | |||
| </Button> | |||
| </div> | |||
| </Card> | |||
| )} | |||
| {userId && ( | |||
| <Card className='!rounded-2xl shadow-sm border-0'> | |||
| <div className='flex items-center mb-2'> | |||
| @@ -395,6 +421,13 @@ const EditUserModal = (props) => { | |||
| formApiRef={formApiRef} | |||
| /> | |||
| <UserVideoChannelBindingModal | |||
| visible={videoBindingModalVisible} | |||
| onCancel={closeVideoBindingModal} | |||
| userId={userId} | |||
| isMobile={isMobile} | |||
| /> | |||
| {/* 添加额度模态框 */} | |||
| <Modal | |||
| centered | |||
| @@ -0,0 +1,156 @@ | |||
| import React, { useEffect, useState } from 'react'; | |||
| import { useTranslation } from 'react-i18next'; | |||
| import { API, showError, showSuccess } from '../../../../helpers'; | |||
| import { | |||
| Button, | |||
| Card, | |||
| Modal, | |||
| Select, | |||
| Spin, | |||
| Typography, | |||
| } from '@douyinfe/semi-ui'; | |||
| const { Text } = Typography; | |||
| const UserVideoChannelBindingModal = ({ | |||
| visible, | |||
| onCancel, | |||
| userId, | |||
| isMobile, | |||
| }) => { | |||
| const { t } = useTranslation(); | |||
| const [loading, setLoading] = useState(false); | |||
| const [saving, setSaving] = useState(false); | |||
| const [data, setData] = useState({ groups: [], families: [] }); | |||
| const [values, setValues] = useState({}); | |||
| useEffect(() => { | |||
| if (!visible || !userId) return; | |||
| const load = async () => { | |||
| setLoading(true); | |||
| try { | |||
| const res = await API.get(`/api/user/${userId}/video-channel-bindings`); | |||
| if (!res.data?.success) { | |||
| showError(res.data?.message || t('操作失败')); | |||
| return; | |||
| } | |||
| const nextData = res.data.data || { groups: [], families: [] }; | |||
| const nextValues = {}; | |||
| nextData.families.forEach((family) => { | |||
| family.bindings.forEach((binding) => { | |||
| nextValues[`${binding.group}:${family.key}`] = | |||
| binding.channel_id || undefined; | |||
| }); | |||
| }); | |||
| setData(nextData); | |||
| setValues(nextValues); | |||
| } catch (error) { | |||
| showError( | |||
| error.response?.data?.message || error.message || t('操作失败'), | |||
| ); | |||
| } finally { | |||
| setLoading(false); | |||
| } | |||
| }; | |||
| load(); | |||
| }, [t, userId, visible]); | |||
| const save = async () => { | |||
| const bindings = []; | |||
| data.families.forEach((family) => { | |||
| family.bindings.forEach((binding) => { | |||
| const channelId = values[`${binding.group}:${family.key}`]; | |||
| if (channelId) { | |||
| bindings.push({ | |||
| group: binding.group, | |||
| family: family.key, | |||
| channel_id: channelId, | |||
| }); | |||
| } | |||
| }); | |||
| }); | |||
| setSaving(true); | |||
| try { | |||
| const res = await API.put(`/api/user/${userId}/video-channel-bindings`, { | |||
| bindings, | |||
| }); | |||
| if (!res.data?.success) { | |||
| showError(res.data?.message || t('操作失败')); | |||
| return; | |||
| } | |||
| showSuccess(t('视频渠道绑定已保存')); | |||
| onCancel(); | |||
| } catch (error) { | |||
| showError( | |||
| error.response?.data?.message || error.message || t('操作失败'), | |||
| ); | |||
| } finally { | |||
| setSaving(false); | |||
| } | |||
| }; | |||
| return ( | |||
| <Modal | |||
| visible={visible} | |||
| onCancel={onCancel} | |||
| title={t('视频渠道绑定')} | |||
| width={isMobile ? '100%' : 760} | |||
| okText={t('保存')} | |||
| cancelText={t('取消')} | |||
| confirmLoading={saving} | |||
| onOk={save} | |||
| > | |||
| <Spin spinning={loading}> | |||
| {data.groups.length === 0 ? ( | |||
| <Text type='tertiary'> | |||
| {t('该用户当前 Token 分组没有可用的视频渠道')} | |||
| </Text> | |||
| ) : ( | |||
| <div className='space-y-4 max-h-[60vh] overflow-y-auto pr-1'> | |||
| {data.groups.map((group) => { | |||
| const availableBindings = data.families.flatMap((family) => { | |||
| const binding = family.bindings.find( | |||
| (item) => item.group === group, | |||
| ); | |||
| return binding ? [{ family, binding }] : []; | |||
| }); | |||
| return ( | |||
| <Card key={group} className='!rounded-xl'> | |||
| <Text strong>{group}</Text> | |||
| <div className='mt-3 space-y-3'> | |||
| {availableBindings.map(({ family, binding }) => { | |||
| const field = `${group}:${family.key}`; | |||
| return ( | |||
| <div key={field}> | |||
| <Text className='block mb-1'>{family.name}</Text> | |||
| <Select | |||
| value={values[field]} | |||
| placeholder={t('未绑定,保存后将恢复自动选择')} | |||
| optionList={binding.candidates.map((channel) => ({ | |||
| label: `${channel.name} (#${channel.id}, type ${channel.type})`, | |||
| value: channel.id, | |||
| }))} | |||
| showClear | |||
| style={{ width: '100%' }} | |||
| onChange={(value) => | |||
| setValues((previous) => ({ | |||
| ...previous, | |||
| [field]: value, | |||
| })) | |||
| } | |||
| /> | |||
| </div> | |||
| ); | |||
| })} | |||
| </div> | |||
| </Card> | |||
| ); | |||
| })} | |||
| </div> | |||
| )} | |||
| </Spin> | |||
| </Modal> | |||
| ); | |||
| }; | |||
| export default UserVideoChannelBindingModal; | |||
| @@ -1367,6 +1367,14 @@ | |||
| "搜索部署名称": "Search deployment name", | |||
| "操作": "Actions", | |||
| "操作失败": "Operation failed", | |||
| "视频渠道绑定": "Video Channel Bindings", | |||
| "视频渠道绑定已保存": "Video channel bindings saved", | |||
| "该用户暂无可配置的 Token 分组;请先创建非 auto 分组的令牌。": "This user has no configurable token groups. Create a non-auto token group first.", | |||
| "该用户当前 Token 分组没有可用的视频渠道": "None of this user's token groups have an available video channel", | |||
| "未绑定,保存后将恢复自动选择": "Unbound; saving restores automatic selection", | |||
| "当前分组没有可用渠道": "No available channels for this group", | |||
| "按该用户 Token 分组配置 Seedance、Kling 等视频路由": "Configure Seedance, Kling, and other video routing by this user's token group", | |||
| "配置视频渠道": "Configure Video Channels", | |||
| "操作失败,请重试": "Operation failed, please retry", | |||
| "操作成功完成!": "Operation completed successfully!", | |||
| "操作暂时被禁用": "Operation temporarily disabled", | |||
| @@ -1351,6 +1351,14 @@ | |||
| "搜索部署名称": "搜索部署名称", | |||
| "操作": "操作", | |||
| "操作失败": "操作失败", | |||
| "视频渠道绑定": "视频渠道绑定", | |||
| "视频渠道绑定已保存": "视频渠道绑定已保存", | |||
| "该用户暂无可配置的 Token 分组;请先创建非 auto 分组的令牌。": "该用户暂无可配置的 Token 分组;请先创建非 auto 分组的令牌。", | |||
| "该用户当前 Token 分组没有可用的视频渠道": "该用户当前 Token 分组没有可用的视频渠道", | |||
| "未绑定,保存后将恢复自动选择": "未绑定,保存后将恢复自动选择", | |||
| "当前分组没有可用渠道": "当前分组没有可用渠道", | |||
| "按该用户 Token 分组配置 Seedance、Kling 等视频路由": "按该用户 Token 分组配置 Seedance、Kling 等视频路由", | |||
| "配置视频渠道": "配置视频渠道", | |||
| "操作失败,请重试": "操作失败,请重试", | |||
| "操作成功完成!": "操作成功完成!", | |||
| "操作暂时被禁用": "操作暂时被禁用", | |||