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 ( | import ( | ||||
| "errors" | "errors" | ||||
| "fmt" | "fmt" | ||||
| "sort" | |||||
| "strings" | "strings" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| @@ -27,7 +28,7 @@ type Token struct { | |||||
| AllowIps *string `json:"allow_ips" gorm:"default:''"` | AllowIps *string `json:"allow_ips" gorm:"default:''"` | ||||
| UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota | UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota | ||||
| Group string `json:"group" gorm:"default:''"` | 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"` | DeletedAt gorm.DeletedAt `gorm:"index"` | ||||
| } | } | ||||
| @@ -64,6 +65,24 @@ func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { | |||||
| return tokens, err | 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 搜索模式。 | // sanitizeLikePattern 校验并清洗用户输入的 LIKE 搜索模式。 | ||||
| // 规则: | // 规则: | ||||
| // 1. 转义 ! 和 _(使用 ! 作为 ESCAPE 字符,兼容 MySQL/PostgreSQL/SQLite) | // 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.GET("/:id/oauth/bindings", controller.GetUserOAuthBindingsByAdmin) | ||||
| adminRoute.DELETE("/:id/oauth/bindings/:provider_id", controller.UnbindCustomOAuthByAdmin) | adminRoute.DELETE("/:id/oauth/bindings/:provider_id", controller.UnbindCustomOAuthByAdmin) | ||||
| adminRoute.DELETE("/:id/bindings/:binding_type", controller.AdminClearUserBinding) | 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.GET("/:id", controller.GetUser) | ||||
| adminRoute.POST("/", controller.CreateUser) | adminRoute.POST("/", controller.CreateUser) | ||||
| adminRoute.POST("/manage", controller.ManageUser) | 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"` | ExpiresIn int `json:"expires_in"` | ||||
| } | } | ||||
| if err := common.DecodeJson(resp.Body, &payload); err != nil { | |||||
| return nil, err | |||||
| } | |||||
| if resp.StatusCode < 200 || resp.StatusCode >= 300 { | if resp.StatusCode < 200 || resp.StatusCode >= 300 { | ||||
| return nil, fmt.Errorf("codex oauth refresh failed: status=%d", resp.StatusCode) | 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 { | if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 { | ||||
| return nil, errors.New("codex oauth refresh response missing fields") | return nil, errors.New("codex oauth refresh response missing fields") | ||||
| @@ -181,12 +181,12 @@ func exchangeCodexAuthorizationCode( | |||||
| RefreshToken string `json:"refresh_token"` | RefreshToken string `json:"refresh_token"` | ||||
| ExpiresIn int `json:"expires_in"` | ExpiresIn int `json:"expires_in"` | ||||
| } | } | ||||
| if err := common.DecodeJson(resp.Body, &payload); err != nil { | |||||
| return nil, err | |||||
| } | |||||
| if resp.StatusCode < 200 || resp.StatusCode >= 300 { | if resp.StatusCode < 200 || resp.StatusCode >= 300 { | ||||
| return nil, fmt.Errorf("codex oauth code exchange failed: status=%d", resp.StatusCode) | 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 { | if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 { | ||||
| return nil, errors.New("codex oauth token response missing fields") | 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() | sqlDB, err := db.DB() | ||||
| require.NoError(t, err) | require.NoError(t, err) | ||||
| sqlDB.SetMaxOpenConns(1) | 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() { | t.Cleanup(func() { | ||||
| _ = sqlDB.Close() | _ = 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]) | action := strings.ToLower(split[1]) | ||||
| if action == "" { | |||||
| return nil | |||||
| } | |||||
| changeParams := &dto.MidjourneyRequest{} | changeParams := &dto.MidjourneyRequest{} | ||||
| changeParams.TaskId = split[0] | 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 { | 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) | 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表示当前对齐的文本串位置 | // 主循环,i表示当前对齐的文本串位置 | ||||
| for i := 0; i <= n-m; { | for i := 0; i <= n-m; { | ||||
| // 检查子串 | // 检查子串 | ||||
| j := 0 | j := 0 | ||||
| for j < m && text[i+j] == pattern[j] { | |||||
| for j < m && textRunes[i+j] == patternRunes[j] { | |||||
| j++ | j++ | ||||
| } | } | ||||
| // 如果完全匹配,返回匹配位置 | // 如果完全匹配,返回匹配位置 | ||||
| @@ -35,7 +44,7 @@ func SundaySearch(text string, pattern string) bool { | |||||
| // 如果还有剩余字符,则检查下一位字符在偏移表中的值 | // 如果还有剩余字符,则检查下一位字符在偏移表中的值 | ||||
| if i+m < n { | if i+m < n { | ||||
| next := rune(text[i+m]) | |||||
| next := textRunes[i+m] | |||||
| if val, ok := offset[next]; ok { | if val, ok := offset[next]; ok { | ||||
| i += val // 存在于偏移表中,进行跳跃 | i += val // 存在于偏移表中,进行跳跃 | ||||
| } else { | } 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 ( | import ( | ||||
| "errors" | "errors" | ||||
| "sort" | |||||
| "strings" | "strings" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| @@ -17,6 +18,10 @@ const ( | |||||
| VideoAssetFamilyKling VideoAssetFamily = "kling" | VideoAssetFamilyKling VideoAssetFamily = "kling" | ||||
| ) | ) | ||||
| func VideoAssetFamilies() []VideoAssetFamily { | |||||
| return []VideoAssetFamily{VideoAssetFamilySeedance, VideoAssetFamilyKling} | |||||
| } | |||||
| func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int { | func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int { | ||||
| switch family { | switch family { | ||||
| case VideoAssetFamilySeedance: | 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 { | func HasVideoAssetChannelForGroupModel(tokenGroup string, modelName string, family VideoAssetFamily) bool { | ||||
| tokenGroup = strings.TrimSpace(tokenGroup) | tokenGroup = strings.TrimSpace(tokenGroup) | ||||
| modelName = strings.TrimSpace(modelName) | 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" | "encoding/json" | ||||
| "fmt" | "fmt" | ||||
| "net/http" | "net/http" | ||||
| "strings" | |||||
| "time" | "time" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| @@ -36,7 +37,7 @@ func SendWebhookNotify(webhookURL string, secret string, data dto.Notify) error | |||||
| // 处理占位符 | // 处理占位符 | ||||
| content := data.Content | content := data.Content | ||||
| for _, value := range data.Values { | for _, value := range data.Values { | ||||
| content = fmt.Sprintf(content, value) | |||||
| content = strings.Replace(content, dto.ContentValueParam, fmt.Sprintf("%v", value), 1) | |||||
| } | } | ||||
| // 构建 webhook 负载 | // 构建 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 { | 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 { | func AutoGroups2JsonString() string { | ||||
| @@ -31,8 +31,12 @@ var Chats = []map[string]string{ | |||||
| } | } | ||||
| func UpdateChatsByJsonString(jsonString string) error { | 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 { | 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 { | 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 { | 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, | IconPlus, | ||||
| } from '@douyinfe/semi-icons'; | } from '@douyinfe/semi-icons'; | ||||
| import UserBindingManagementModal from './UserBindingManagementModal'; | import UserBindingManagementModal from './UserBindingManagementModal'; | ||||
| import UserVideoChannelBindingModal from './UserVideoChannelBindingModal'; | |||||
| import UserModelRateLimitSection from './UserModelRateLimitSection'; | import UserModelRateLimitSection from './UserModelRateLimitSection'; | ||||
| const { Text, Title } = Typography; | const { Text, Title } = Typography; | ||||
| @@ -70,6 +71,7 @@ const EditUserModal = (props) => { | |||||
| const isMobile = useIsMobile(); | const isMobile = useIsMobile(); | ||||
| const [groupOptions, setGroupOptions] = useState([]); | const [groupOptions, setGroupOptions] = useState([]); | ||||
| const [bindingModalVisible, setBindingModalVisible] = useState(false); | const [bindingModalVisible, setBindingModalVisible] = useState(false); | ||||
| const [videoBindingModalVisible, setVideoBindingModalVisible] = useState(false); | |||||
| const formApiRef = useRef(null); | const formApiRef = useRef(null); | ||||
| const [isSyncedUser, setIsSyncedUser] = useState(false); | const [isSyncedUser, setIsSyncedUser] = useState(false); | ||||
| @@ -122,6 +124,7 @@ const EditUserModal = (props) => { | |||||
| loadUser(); | loadUser(); | ||||
| if (userId) fetchGroups(); | if (userId) fetchGroups(); | ||||
| setBindingModalVisible(false); | setBindingModalVisible(false); | ||||
| setVideoBindingModalVisible(false); | |||||
| }, [props.editingUser.id]); | }, [props.editingUser.id]); | ||||
| const openBindingModal = () => { | const openBindingModal = () => { | ||||
| @@ -132,6 +135,9 @@ const EditUserModal = (props) => { | |||||
| setBindingModalVisible(false); | setBindingModalVisible(false); | ||||
| }; | }; | ||||
| const openVideoBindingModal = () => setVideoBindingModalVisible(true); | |||||
| const closeVideoBindingModal = () => setVideoBindingModalVisible(false); | |||||
| /* ----------------------- submit ----------------------- */ | /* ----------------------- submit ----------------------- */ | ||||
| const submit = async (values) => { | const submit = async (values) => { | ||||
| setLoading(true); | setLoading(true); | ||||
| @@ -284,6 +290,26 @@ const EditUserModal = (props) => { | |||||
| </Card> | </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 && ( | {userId && ( | ||||
| <Card className='!rounded-2xl shadow-sm border-0'> | <Card className='!rounded-2xl shadow-sm border-0'> | ||||
| <div className='flex items-center mb-2'> | <div className='flex items-center mb-2'> | ||||
| @@ -395,6 +421,13 @@ const EditUserModal = (props) => { | |||||
| formApiRef={formApiRef} | formApiRef={formApiRef} | ||||
| /> | /> | ||||
| <UserVideoChannelBindingModal | |||||
| visible={videoBindingModalVisible} | |||||
| onCancel={closeVideoBindingModal} | |||||
| userId={userId} | |||||
| isMobile={isMobile} | |||||
| /> | |||||
| {/* 添加额度模态框 */} | {/* 添加额度模态框 */} | ||||
| <Modal | <Modal | ||||
| centered | 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", | "搜索部署名称": "Search deployment name", | ||||
| "操作": "Actions", | "操作": "Actions", | ||||
| "操作失败": "Operation failed", | "操作失败": "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 failed, please retry", | ||||
| "操作成功完成!": "Operation completed successfully!", | "操作成功完成!": "Operation completed successfully!", | ||||
| "操作暂时被禁用": "Operation temporarily disabled", | "操作暂时被禁用": "Operation temporarily disabled", | ||||
| @@ -1351,6 +1351,14 @@ | |||||
| "搜索部署名称": "搜索部署名称", | "搜索部署名称": "搜索部署名称", | ||||
| "操作": "操作", | "操作": "操作", | ||||
| "操作失败": "操作失败", | "操作失败": "操作失败", | ||||
| "视频渠道绑定": "视频渠道绑定", | |||||
| "视频渠道绑定已保存": "视频渠道绑定已保存", | |||||
| "该用户暂无可配置的 Token 分组;请先创建非 auto 分组的令牌。": "该用户暂无可配置的 Token 分组;请先创建非 auto 分组的令牌。", | |||||
| "该用户当前 Token 分组没有可用的视频渠道": "该用户当前 Token 分组没有可用的视频渠道", | |||||
| "未绑定,保存后将恢复自动选择": "未绑定,保存后将恢复自动选择", | |||||
| "当前分组没有可用渠道": "当前分组没有可用渠道", | |||||
| "按该用户 Token 分组配置 Seedance、Kling 等视频路由": "按该用户 Token 分组配置 Seedance、Kling 等视频路由", | |||||
| "配置视频渠道": "配置视频渠道", | |||||
| "操作失败,请重试": "操作失败,请重试", | "操作失败,请重试": "操作失败,请重试", | ||||
| "操作成功完成!": "操作成功完成!", | "操作成功完成!": "操作成功完成!", | ||||
| "操作暂时被禁用": "操作暂时被禁用", | "操作暂时被禁用": "操作暂时被禁用", | ||||