diff --git a/controller/channel_affinity_cache_test.go b/controller/channel_affinity_cache_test.go new file mode 100644 index 0000000..e946c2b --- /dev/null +++ b/controller/channel_affinity_cache_test.go @@ -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`) +} diff --git a/controller/channel_billing_helpers_test.go b/controller/channel_billing_helpers_test.go new file mode 100644 index 0000000..ff7ebe0 --- /dev/null +++ b/controller/channel_billing_helpers_test.go @@ -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) +} diff --git a/controller/deployment_test.go b/controller/deployment_test.go new file mode 100644 index 0000000..69668f9 --- /dev/null +++ b/controller/deployment_test.go @@ -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) +} diff --git a/controller/helpers_test.go b/controller/helpers_test.go new file mode 100644 index 0000000..abd5f31 --- /dev/null +++ b/controller/helpers_test.go @@ -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")) +} diff --git a/controller/option_test.go b/controller/option_test.go new file mode 100644 index 0000000..0c6168d --- /dev/null +++ b/controller/option_test.go @@ -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") +} diff --git a/controller/payment_webhook_availability_test.go b/controller/payment_webhook_availability_test.go new file mode 100644 index 0000000..7f635ff --- /dev/null +++ b/controller/payment_webhook_availability_test.go @@ -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()) +} diff --git a/controller/ratio_config_test.go b/controller/ratio_config_test.go new file mode 100644 index 0000000..30280be --- /dev/null +++ b/controller/ratio_config_test.go @@ -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()) +} diff --git a/controller/user_video_channel_binding.go b/controller/user_video_channel_binding.go new file mode 100644 index 0000000..dea8893 --- /dev/null +++ b/controller/user_video_channel_binding.go @@ -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": ""}) +} diff --git a/controller/user_video_channel_binding_test.go b/controller/user_video_channel_binding_test.go new file mode 100644 index 0000000..4d5fabd --- /dev/null +++ b/controller/user_video_channel_binding_test.go @@ -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) + } +} diff --git a/controller/video_proxy_gemini_test.go b/controller/video_proxy_gemini_test.go new file mode 100644 index 0000000..2fba96a --- /dev/null +++ b/controller/video_proxy_gemini_test.go @@ -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", "")) +} diff --git a/docs/testing/2026-07-24-admin-video-channel-binding.md b/docs/testing/2026-07-24-admin-video-channel-binding.md new file mode 100644 index 0000000..8b5a70c --- /dev/null +++ b/docs/testing/2026-07-24-admin-video-channel-binding.md @@ -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` 未导出该符号。新增中英文键已手动同步，生产构建已验证通过。 diff --git a/middleware/gzip_test.go b/middleware/gzip_test.go new file mode 100644 index 0000000..c3dee4a --- /dev/null +++ b/middleware/gzip_test.go @@ -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) +} diff --git a/middleware/http_headers_test.go b/middleware/http_headers_test.go new file mode 100644 index 0000000..37367f4 --- /dev/null +++ b/middleware/http_headers_test.go @@ -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)) +} diff --git a/middleware/recover_test.go b/middleware/recover_test.go new file mode 100644 index 0000000..63da656 --- /dev/null +++ b/middleware/recover_test.go @@ -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) +} diff --git a/model/custom_oauth_provider_test.go b/model/custom_oauth_provider_test.go new file mode 100644 index 0000000..87aa042 --- /dev/null +++ b/model/custom_oauth_provider_test.go @@ -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") +} diff --git a/model/option_map_test.go b/model/option_map_test.go new file mode 100644 index 0000000..da3d3dd --- /dev/null +++ b/model/option_map_test.go @@ -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()) +} diff --git a/model/passkey_test.go b/model/passkey_test.go new file mode 100644 index 0000000..05b1148 --- /dev/null +++ b/model/passkey_test.go @@ -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) +} diff --git a/model/prefill_group_test.go b/model/prefill_group_test.go new file mode 100644 index 0000000..051e8f7 --- /dev/null +++ b/model/prefill_group_test.go @@ -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) +} diff --git a/model/pricing_default_test.go b/model/pricing_default_test.go new file mode 100644 index 0000000..924cdde --- /dev/null +++ b/model/pricing_default_test.go @@ -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")) +} diff --git a/model/token.go b/model/token.go index 2f318f8..5999e4e 100644 --- a/model/token.go +++ b/model/token.go @@ -3,6 +3,7 @@ package model import ( "errors" "fmt" + "sort" "strings" "github.com/QuantumNous/new-api/common" @@ -27,7 +28,7 @@ type Token struct { AllowIps *string `json:"allow_ips" gorm:"default:''"` UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota Group string `json:"group" gorm:"default:''"` - CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效 + CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效 DeletedAt gorm.DeletedAt `gorm:"index"` } @@ -64,6 +65,24 @@ func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { return tokens, err } +// GetUserConcreteTokenGroups returns the groups an administrator can bind. +// The auto group resolves at request time and cannot be bound directly. +func GetUserConcreteTokenGroups(userId int) ([]string, error) { + var rawGroups []string + if err := DB.Model(&Token{}).Where("user_id = ?", userId).Distinct().Pluck("group", &rawGroups).Error; err != nil { + return nil, err + } + groups := make([]string, 0, len(rawGroups)) + for _, group := range rawGroups { + group = strings.TrimSpace(group) + if group != "" && group != "auto" { + groups = append(groups, group) + } + } + sort.Strings(groups) + return groups, nil +} + // sanitizeLikePattern 校验并清洗用户输入的 LIKE 搜索模式。 // 规则： // 1. 转义 ! 和 _（使用 ! 作为 ESCAPE 字符，兼容 MySQL/PostgreSQL/SQLite） diff --git a/model/token_helpers_test.go b/model/token_helpers_test.go new file mode 100644 index 0000000..1c815e7 --- /dev/null +++ b/model/token_helpers_test.go @@ -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()) +} diff --git a/model/twofa_test.go b/model/twofa_test.go new file mode 100644 index 0000000..bd68830 --- /dev/null +++ b/model/twofa_test.go @@ -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()) +} diff --git a/model/user_cache_test.go b/model/user_cache_test.go new file mode 100644 index 0000000..91775c0 --- /dev/null +++ b/model/user_cache_test.go @@ -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)) +} diff --git a/model/utils_test.go b/model/utils_test.go new file mode 100644 index 0000000..61865c2 --- /dev/null +++ b/model/utils_test.go @@ -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)) +} diff --git a/router/api-router.go b/router/api-router.go index 67e536c..98d0987 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -136,6 +136,8 @@ func SetApiRouter(router *gin.Engine) { adminRoute.GET("/:id/oauth/bindings", controller.GetUserOAuthBindingsByAdmin) adminRoute.DELETE("/:id/oauth/bindings/:provider_id", controller.UnbindCustomOAuthByAdmin) adminRoute.DELETE("/:id/bindings/:binding_type", controller.AdminClearUserBinding) + adminRoute.GET("/:id/video-channel-bindings", controller.GetUserVideoChannelBindings) + adminRoute.PUT("/:id/video-channel-bindings", controller.SetUserVideoChannelBindings) adminRoute.GET("/:id", controller.GetUser) adminRoute.POST("/", controller.CreateUser) adminRoute.POST("/manage", controller.ManageUser) diff --git a/service/audio_test.go b/service/audio_test.go new file mode 100644 index 0000000..ec4664e --- /dev/null +++ b/service/audio_test.go @@ -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) +} diff --git a/service/channel_affinity_helpers_test.go b/service/channel_affinity_helpers_test.go new file mode 100644 index 0000000..968c04a --- /dev/null +++ b/service/channel_affinity_helpers_test.go @@ -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)) +} diff --git a/service/channel_test.go b/service/channel_test.go new file mode 100644 index 0000000..a2caec2 --- /dev/null +++ b/service/channel_test.go @@ -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)) +} diff --git a/service/codex_oauth.go b/service/codex_oauth.go index 33ef1d6..9de84e1 100644 --- a/service/codex_oauth.go +++ b/service/codex_oauth.go @@ -120,12 +120,12 @@ func refreshCodexOAuthToken( ExpiresIn int `json:"expires_in"` } - if err := common.DecodeJson(resp.Body, &payload); err != nil { - return nil, err - } if resp.StatusCode < 200 || resp.StatusCode >= 300 { return nil, fmt.Errorf("codex oauth refresh failed: status=%d", resp.StatusCode) } + if err := common.DecodeJson(resp.Body, &payload); err != nil { + return nil, err + } if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 { return nil, errors.New("codex oauth refresh response missing fields") @@ -181,12 +181,12 @@ func exchangeCodexAuthorizationCode( RefreshToken string `json:"refresh_token"` ExpiresIn int `json:"expires_in"` } - if err := common.DecodeJson(resp.Body, &payload); err != nil { - return nil, err - } if resp.StatusCode < 200 || resp.StatusCode >= 300 { return nil, fmt.Errorf("codex oauth code exchange failed: status=%d", resp.StatusCode) } + if err := common.DecodeJson(resp.Body, &payload); err != nil { + return nil, err + } if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 { return nil, errors.New("codex oauth token response missing fields") } diff --git a/service/codex_oauth_test.go b/service/codex_oauth_test.go new file mode 100644 index 0000000..23f6a8e --- /dev/null +++ b/service/codex_oauth_test.go @@ -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" +} diff --git a/service/convert_helpers_test.go b/service/convert_helpers_test.go new file mode 100644 index 0000000..a0fa4cb --- /dev/null +++ b/service/convert_helpers_test.go @@ -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")) +} diff --git a/service/doubao_asset_channel_test.go b/service/doubao_asset_channel_test.go index 0a14e2f..da690c5 100644 --- a/service/doubao_asset_channel_test.go +++ b/service/doubao_asset_channel_test.go @@ -42,7 +42,7 @@ func setupDoubaoAssetChannelDB(t *testing.T) *gorm.DB { sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) - require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{})) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.Token{})) t.Cleanup(func() { _ = sqlDB.Close() diff --git a/service/file_service_test.go b/service/file_service_test.go new file mode 100644 index 0000000..a34edd1 --- /dev/null +++ b/service/file_service_test.go @@ -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")) +} diff --git a/service/group_test.go b/service/group_test.go new file mode 100644 index 0000000..de912fe --- /dev/null +++ b/service/group_test.go @@ -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")) +} diff --git a/service/http_client_test.go b/service/http_client_test.go new file mode 100644 index 0000000..c42a965 --- /dev/null +++ b/service/http_client_test.go @@ -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) +} diff --git a/service/image_test.go b/service/image_test.go new file mode 100644 index 0000000..cd043fa --- /dev/null +++ b/service/image_test.go @@ -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) +} diff --git a/service/midjourney.go b/service/midjourney.go index bdb0fe5..7515df9 100644 --- a/service/midjourney.go +++ b/service/midjourney.go @@ -140,6 +140,9 @@ func ConvertSimpleChangeParams(content string) *dto.MidjourneyRequest { } action := strings.ToLower(split[1]) + if action == "" { + return nil + } changeParams := &dto.MidjourneyRequest{} changeParams.TaskId = split[0] diff --git a/service/midjourney_helpers_test.go b/service/midjourney_helpers_test.go new file mode 100644 index 0000000..c0b4152 --- /dev/null +++ b/service/midjourney_helpers_test.go @@ -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")) +} diff --git a/service/openaicompat/policy_test.go b/service/openaicompat/policy_test.go new file mode 100644 index 0000000..6728c8a --- /dev/null +++ b/service/openaicompat/policy_test.go @@ -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")) +} diff --git a/service/passkey/service_test.go b/service/passkey/service_test.go new file mode 100644 index 0000000..c201973 --- /dev/null +++ b/service/passkey/service_test.go @@ -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")) +} diff --git a/service/sensitive_test.go b/service/sensitive_test.go new file mode 100644 index 0000000..57f9d60 --- /dev/null +++ b/service/sensitive_test.go @@ -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) +} diff --git a/service/str.go b/service/str.go index 61054bd..542cad0 100644 --- a/service/str.go +++ b/service/str.go @@ -12,20 +12,29 @@ import ( ) func SundaySearch(text string, pattern string) bool { + textRunes := []rune(text) + patternRunes := []rune(pattern) + if len(patternRunes) == 0 { + return true + } + if len(patternRunes) > len(textRunes) { + return false + } + // 计算偏移表 offset := make(map[rune]int) - for i, c := range pattern { - offset[c] = len(pattern) - i + for i, c := range patternRunes { + offset[c] = len(patternRunes) - i } // 文本串长度和模式串长度 - n, m := len(text), len(pattern) + n, m := len(textRunes), len(patternRunes) // 主循环，i表示当前对齐的文本串位置 for i := 0; i <= n-m; { // 检查子串 j := 0 - for j < m && text[i+j] == pattern[j] { + for j < m && textRunes[i+j] == patternRunes[j] { j++ } // 如果完全匹配，返回匹配位置 @@ -35,7 +44,7 @@ func SundaySearch(text string, pattern string) bool { // 如果还有剩余字符，则检查下一位字符在偏移表中的值 if i+m < n { - next := rune(text[i+m]) + next := textRunes[i+m] if val, ok := offset[next]; ok { i += val // 存在于偏移表中，进行跳跃 } else { diff --git a/service/str_test.go b/service/str_test.go new file mode 100644 index 0000000..c7debaa --- /dev/null +++ b/service/str_test.go @@ -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"})) +} diff --git a/service/token_estimator_test.go b/service/token_estimator_test.go new file mode 100644 index 0000000..dd1656c --- /dev/null +++ b/service/token_estimator_test.go @@ -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")) +} diff --git a/service/usage_helpr_test.go b/service/usage_helpr_test.go new file mode 100644 index 0000000..a161d29 --- /dev/null +++ b/service/usage_helpr_test.go @@ -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})) +} diff --git a/service/video_asset_channel.go b/service/video_asset_channel.go index e3d6019..42fd84b 100644 --- a/service/video_asset_channel.go +++ b/service/video_asset_channel.go @@ -2,6 +2,7 @@ package service import ( "errors" + "sort" "strings" "github.com/QuantumNous/new-api/common" @@ -17,6 +18,10 @@ const ( VideoAssetFamilyKling VideoAssetFamily = "kling" ) +func VideoAssetFamilies() []VideoAssetFamily { + return []VideoAssetFamily{VideoAssetFamilySeedance, VideoAssetFamilyKling} +} + func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int { switch family { case VideoAssetFamilySeedance: @@ -107,6 +112,55 @@ func BindVideoAssetChannel(userId int, tokenGroup string, channel *model.Channel }) } +func ClearVideoAssetChannelBinding(userId int, tokenGroup string, family VideoAssetFamily) error { + channelTypes := VideoAssetChannelTypesForFamily(family) + if len(channelTypes) == 0 { + return nil + } + return model.DB.Transaction(func(tx *gorm.DB) error { + return model.DeleteUserAssetChannelsByTypesWithTx(tx, userId, channelTypes, tokenGroup) + }) +} + +func GetVideoAssetChannelCandidates(tokenGroup string, family VideoAssetFamily) ([]*model.Channel, error) { + channelTypes := VideoAssetChannelTypesForFamily(family) + if len(channelTypes) == 0 { + return nil, nil + } + candidates := make(map[int]*model.Channel) + for _, channelType := range channelTypes { + for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { + channels, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) + if err != nil { + return nil, err + } + for _, candidate := range channels { + if candidate == nil || candidate.Status != common.ChannelStatusEnabled || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { + continue + } + channel, err := model.GetChannelById(candidate.Id, true) + if err != nil { + return nil, err + } + if channel.Status == common.ChannelStatusEnabled && strings.TrimSpace(channel.Key) != "" && MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) { + channel.Key = "" + channel.Keys = nil + candidates[channel.Id] = channel + } + } + if len(channels) < DoubaoAssetChannelPageSize { + break + } + } + } + result := make([]*model.Channel, 0, len(candidates)) + for _, channel := range candidates { + result = append(result, channel) + } + sort.Slice(result, func(i, j int) bool { return result[i].Id < result[j].Id }) + return result, nil +} + func HasVideoAssetChannelForGroupModel(tokenGroup string, modelName string, family VideoAssetFamily) bool { tokenGroup = strings.TrimSpace(tokenGroup) modelName = strings.TrimSpace(modelName) diff --git a/service/video_asset_channel_admin_test.go b/service/video_asset_channel_admin_test.go new file mode 100644 index 0000000..6ee547b --- /dev/null +++ b/service/video_asset_channel_admin_test.go @@ -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) +} diff --git a/service/violation_fee_test.go b/service/violation_fee_test.go new file mode 100644 index 0000000..ae05675 --- /dev/null +++ b/service/violation_fee_test.go @@ -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)) +} diff --git a/service/webhook.go b/service/webhook.go index bab8842..8ce8236 100644 --- a/service/webhook.go +++ b/service/webhook.go @@ -8,6 +8,7 @@ import ( "encoding/json" "fmt" "net/http" + "strings" "time" "github.com/QuantumNous/new-api/common" @@ -36,7 +37,7 @@ func SendWebhookNotify(webhookURL string, secret string, data dto.Notify) error // 处理占位符 content := data.Content for _, value := range data.Values { - content = fmt.Sprintf(content, value) + content = strings.Replace(content, dto.ContentValueParam, fmt.Sprintf("%v", value), 1) } // 构建 webhook 负载 diff --git a/service/webhook_test.go b/service/webhook_test.go new file mode 100644 index 0000000..00c1d0a --- /dev/null +++ b/service/webhook_test.go @@ -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) +} diff --git a/setting/auto_group.go b/setting/auto_group.go index 9261286..e0cc92d 100644 --- a/setting/auto_group.go +++ b/setting/auto_group.go @@ -20,8 +20,12 @@ func ContainsAutoGroup(group string) bool { } func UpdateAutoGroupsByJsonString(jsonString string) error { - autoGroups = make([]string, 0) - return common.Unmarshal([]byte(jsonString), &autoGroups) + var groups []string + if err := common.Unmarshal([]byte(jsonString), &groups); err != nil { + return err + } + autoGroups = groups + return nil } func AutoGroups2JsonString() string { diff --git a/setting/chat.go b/setting/chat.go index 9be0826..4083003 100644 --- a/setting/chat.go +++ b/setting/chat.go @@ -31,8 +31,12 @@ var Chats = []map[string]string{ } func UpdateChatsByJsonString(jsonString string) error { - Chats = make([]map[string]string, 0) - return json.Unmarshal([]byte(jsonString), &Chats) + var chats []map[string]string + if err := json.Unmarshal([]byte(jsonString), &chats); err != nil { + return err + } + Chats = chats + return nil } func Chats2JsonString() string { diff --git a/setting/collection_updates_test.go b/setting/collection_updates_test.go new file mode 100644 index 0000000..d165930 --- /dev/null +++ b/setting/collection_updates_test.go @@ -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) +} diff --git a/setting/operation_setting/payment_setting_old.go b/setting/operation_setting/payment_setting_old.go index d34b6f0..bc225b1 100644 --- a/setting/operation_setting/payment_setting_old.go +++ b/setting/operation_setting/payment_setting_old.go @@ -37,8 +37,12 @@ var PayMethods = []map[string]string{ } func UpdatePayMethodsByJsonString(jsonString string) error { - PayMethods = make([]map[string]string, 0) - return common.Unmarshal([]byte(jsonString), &PayMethods) + var methods []map[string]string + if err := common.Unmarshal([]byte(jsonString), &methods); err != nil { + return err + } + PayMethods = methods + return nil } func PayMethods2JsonString() string { diff --git a/setting/operation_setting/payment_setting_old_test.go b/setting/operation_setting/payment_setting_old_test.go new file mode 100644 index 0000000..5a6a9cd --- /dev/null +++ b/setting/operation_setting/payment_setting_old_test.go @@ -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) +} diff --git a/web/src/components/table/users/modals/EditUserModal.jsx b/web/src/components/table/users/modals/EditUserModal.jsx index 2bda2f5..91dd569 100644 --- a/web/src/components/table/users/modals/EditUserModal.jsx +++ b/web/src/components/table/users/modals/EditUserModal.jsx @@ -56,6 +56,7 @@ import { IconPlus, } from '@douyinfe/semi-icons'; import UserBindingManagementModal from './UserBindingManagementModal'; +import UserVideoChannelBindingModal from './UserVideoChannelBindingModal'; import UserModelRateLimitSection from './UserModelRateLimitSection'; const { Text, Title } = Typography; @@ -70,6 +71,7 @@ const EditUserModal = (props) => { const isMobile = useIsMobile(); const [groupOptions, setGroupOptions] = useState([]); const [bindingModalVisible, setBindingModalVisible] = useState(false); + const [videoBindingModalVisible, setVideoBindingModalVisible] = useState(false); const formApiRef = useRef(null); const [isSyncedUser, setIsSyncedUser] = useState(false); @@ -122,6 +124,7 @@ const EditUserModal = (props) => { loadUser(); if (userId) fetchGroups(); setBindingModalVisible(false); + setVideoBindingModalVisible(false); }, [props.editingUser.id]); const openBindingModal = () => { @@ -132,6 +135,9 @@ const EditUserModal = (props) => { setBindingModalVisible(false); }; + const openVideoBindingModal = () => setVideoBindingModalVisible(true); + const closeVideoBindingModal = () => setVideoBindingModalVisible(false); + /* ----------------------- submit ----------------------- */ const submit = async (values) => { setLoading(true); @@ -284,6 +290,26 @@ const EditUserModal = (props) => { {/* 权限设置 */} + {userId && ( + +
+
+ + + +
+ {t('视频渠道绑定')} +
+ {t('按该用户 Token 分组配置 Seedance、Kling 等视频路由')} +
+
+
+ +
+ + )} {userId && (
@@ -395,6 +421,13 @@ const EditUserModal = (props) => { formApiRef={formApiRef} /> + + {/* 添加额度模态框 */} { + 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 ( + + + {data.groups.length === 0 ? ( + + {t('该用户当前 Token 分组没有可用的视频渠道')} + + ) : ( +
+ {data.groups.map((group) => { + const availableBindings = data.families.flatMap((family) => { + const binding = family.bindings.find( + (item) => item.group === group, + ); + return binding ? [{ family, binding }] : []; + }); + return ( + + {group} +
+ {availableBindings.map(({ family, binding }) => { + const field = `${group}:${family.key}`; + return ( +
+ {family.name} +