Add unit tests for group ratio, cross-group retry, and price calculation: - Group ratio CRUD and special ratio override (12 tests) - HandleGroupRatio with auto-group and special ratios (5 tests) - ModelPriceHelper with group ratio in price/ratio modes (3 tests) - Channel selection priority, groups, weighted LB (6 tests) - RetryParam and AutoGroup service logic (7 tests) Also restore the group pricing card in model detail SideSheet.master
| @@ -0,0 +1,187 @@ | |||||
| package model | |||||
| import ( | |||||
| "testing" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/glebarez/sqlite" | |||||
| "github.com/stretchr/testify/assert" | |||||
| "github.com/stretchr/testify/require" | |||||
| "gorm.io/gorm" | |||||
| ) | |||||
| func setupChannelSelectDB(t *testing.T) *gorm.DB { | |||||
| t.Helper() | |||||
| db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) | |||||
| require.NoError(t, err) | |||||
| sqlDB, err := db.DB() | |||||
| require.NoError(t, err) | |||||
| sqlDB.SetMaxOpenConns(1) | |||||
| origDB := DB | |||||
| origUsingSQLite := common.UsingSQLite | |||||
| origRedisEnabled := common.RedisEnabled | |||||
| origMemoryCache := common.MemoryCacheEnabled | |||||
| origGroupCol := commonGroupCol | |||||
| origTrueVal := commonTrueVal | |||||
| DB = db | |||||
| common.UsingSQLite = true | |||||
| common.RedisEnabled = false | |||||
| common.MemoryCacheEnabled = false | |||||
| // Set column quoting for SQLite | |||||
| commonGroupCol = "`group`" | |||||
| commonTrueVal = "1" | |||||
| require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{})) | |||||
| t.Cleanup(func() { | |||||
| DB = origDB | |||||
| common.UsingSQLite = origUsingSQLite | |||||
| common.RedisEnabled = origRedisEnabled | |||||
| common.MemoryCacheEnabled = origMemoryCache | |||||
| commonGroupCol = origGroupCol | |||||
| commonTrueVal = origTrueVal | |||||
| sqlDB.Close() | |||||
| }) | |||||
| return db | |||||
| } | |||||
| func int64Ptr(v int64) *int64 { return &v } | |||||
| func intPtr(v int) *int { return &v } | |||||
| func uintPtr(v uint) *uint { return &v } | |||||
| func createChannelForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, priority int64) { | |||||
| t.Helper() | |||||
| ch := Channel{ | |||||
| Id: id, | |||||
| Type: 1, | |||||
| Key: "test-key-" + name, | |||||
| Status: common.ChannelStatusEnabled, | |||||
| Name: name, | |||||
| Group: group, | |||||
| Models: models, | |||||
| Priority: int64Ptr(priority), | |||||
| Weight: uintPtr(10), | |||||
| AutoBan: intPtr(1), | |||||
| CreatedTime: 1, | |||||
| } | |||||
| require.NoError(t, db.Create(&ch).Error) | |||||
| } | |||||
| func createAbilityForTest(t *testing.T, db *gorm.DB, group, modelName string, channelId int, priority int64, enabled bool) { | |||||
| t.Helper() | |||||
| ability := Ability{ | |||||
| Group: group, | |||||
| Model: modelName, | |||||
| ChannelId: channelId, | |||||
| Enabled: enabled, | |||||
| Priority: int64Ptr(priority), | |||||
| Weight: 10, | |||||
| } | |||||
| require.NoError(t, db.Create(&ability).Error) | |||||
| } | |||||
| // --- GetChannel (DB path) tests --- | |||||
| func TestGetChannel_BasicSelection(t *testing.T) { | |||||
| db := setupChannelSelectDB(t) | |||||
| createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true) | |||||
| ch, err := GetChannel("default", "gpt-4", 0) | |||||
| require.NoError(t, err) | |||||
| assert.NotNil(t, ch) | |||||
| assert.Equal(t, 1, ch.Id) | |||||
| } | |||||
| func TestGetChannel_NoChannelForModel(t *testing.T) { | |||||
| db := setupChannelSelectDB(t) | |||||
| createChannelForTest(t, db, 1, "ch1", "default", "gpt-3.5", 10) | |||||
| createAbilityForTest(t, db, "default", "gpt-3.5", 1, 10, true) | |||||
| ch, err := GetChannel("default", "gpt-4", 0) | |||||
| assert.NoError(t, err) | |||||
| assert.Nil(t, ch) | |||||
| } | |||||
| func TestGetChannel_PrioritySelection(t *testing.T) { | |||||
| db := setupChannelSelectDB(t) | |||||
| // Two channels with different priorities | |||||
| createChannelForTest(t, db, 1, "ch-low", "default", "gpt-4", 0) | |||||
| createChannelForTest(t, db, 2, "ch-high", "default", "gpt-4", 10) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 1, 0, true) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true) | |||||
| // Retry=0 should get highest priority (10) | |||||
| ch, err := GetChannel("default", "gpt-4", 0) | |||||
| require.NoError(t, err) | |||||
| assert.NotNil(t, ch) | |||||
| assert.Equal(t, 2, ch.Id) // ch-high has priority 10 | |||||
| // Retry=1 should get next priority (0) | |||||
| ch, err = GetChannel("default", "gpt-4", 1) | |||||
| require.NoError(t, err) | |||||
| assert.NotNil(t, ch) | |||||
| assert.Equal(t, 1, ch.Id) // ch-low has priority 0 | |||||
| } | |||||
| func TestGetChannel_DifferentGroups(t *testing.T) { | |||||
| db := setupChannelSelectDB(t) | |||||
| // Channel in "default" group | |||||
| createChannelForTest(t, db, 1, "ch-default", "default", "gpt-4", 10) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true) | |||||
| // Channel in "vip" group | |||||
| createChannelForTest(t, db, 2, "ch-vip", "vip", "gpt-4", 10) | |||||
| createAbilityForTest(t, db, "vip", "gpt-4", 2, 10, true) | |||||
| // Query default group | |||||
| ch, err := GetChannel("default", "gpt-4", 0) | |||||
| require.NoError(t, err) | |||||
| assert.Equal(t, 1, ch.Id) | |||||
| // Query vip group | |||||
| ch, err = GetChannel("vip", "gpt-4", 0) | |||||
| require.NoError(t, err) | |||||
| assert.Equal(t, 2, ch.Id) | |||||
| } | |||||
| func TestGetChannel_DisabledAbility(t *testing.T) { | |||||
| db := setupChannelSelectDB(t) | |||||
| createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 1, 10, false) // disabled | |||||
| ch, err := GetChannel("default", "gpt-4", 0) | |||||
| assert.NoError(t, err) | |||||
| assert.Nil(t, ch) // no enabled ability | |||||
| } | |||||
| func TestGetChannel_MultipleChannelsWeighted(t *testing.T) { | |||||
| db := setupChannelSelectDB(t) | |||||
| createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10) | |||||
| createChannelForTest(t, db, 2, "ch2", "default", "gpt-4", 10) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true) | |||||
| createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true) | |||||
| // Run multiple times to verify both channels can be selected | |||||
| found := make(map[int]bool) | |||||
| for i := 0; i < 50; i++ { | |||||
| ch, err := GetChannel("default", "gpt-4", 0) | |||||
| require.NoError(t, err) | |||||
| require.NotNil(t, ch) | |||||
| found[ch.Id] = true | |||||
| } | |||||
| // With equal weights, both should appear at least once in 50 tries | |||||
| assert.True(t, found[1], "channel 1 should be selected at least once") | |||||
| assert.True(t, found[2], "channel 2 should be selected at least once") | |||||
| } | |||||
| @@ -10,6 +10,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/dto" | "github.com/QuantumNous/new-api/dto" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | relaycommon "github.com/QuantumNous/new-api/relay/common" | ||||
| "github.com/QuantumNous/new-api/setting/operation_setting" | |||||
| "github.com/QuantumNous/new-api/setting/ratio_setting" | "github.com/QuantumNous/new-api/setting/ratio_setting" | ||||
| "github.com/QuantumNous/new-api/types" | "github.com/QuantumNous/new-api/types" | ||||
| "github.com/gin-gonic/gin" | "github.com/gin-gonic/gin" | ||||
| @@ -148,3 +149,194 @@ func TestUpdatePriceDataForChannelPricingOnlyAppliesUserChannelRatio(t *testing. | |||||
| assert.Equal(t, 800, info.PriceData.QuotaToPreConsume) | assert.Equal(t, 800, info.PriceData.QuotaToPreConsume) | ||||
| assert.Equal(t, 0.5, info.PriceData.ModelPrice) | assert.Equal(t, 0.5, info.PriceData.ModelPrice) | ||||
| } | } | ||||
| // --- HandleGroupRatio tests --- | |||||
| func TestHandleGroupRatio_DefaultGroup(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| c := buildTestContext(t) | |||||
| info := &relaycommon.RelayInfo{ | |||||
| UsingGroup: "default", | |||||
| UserGroup: "default", | |||||
| } | |||||
| groupRatioInfo := HandleGroupRatio(c, info) | |||||
| assert.Equal(t, 1.0, groupRatioInfo.GroupRatio) | |||||
| assert.False(t, groupRatioInfo.HasSpecialRatio) | |||||
| assert.Equal(t, -1.0, groupRatioInfo.GroupSpecialRatio) | |||||
| } | |||||
| func TestHandleGroupRatio_CustomGroupRatio(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"premium":0.7}`)) | |||||
| c := buildTestContext(t) | |||||
| info := &relaycommon.RelayInfo{ | |||||
| UsingGroup: "premium", | |||||
| UserGroup: "default", | |||||
| } | |||||
| groupRatioInfo := HandleGroupRatio(c, info) | |||||
| assert.Equal(t, 0.7, groupRatioInfo.GroupRatio) | |||||
| assert.False(t, groupRatioInfo.HasSpecialRatio) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||||
| } | |||||
| func TestHandleGroupRatio_SpecialGroupGroupRatio(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{"vip":{"premium":0.8}}`)) | |||||
| c := buildTestContext(t) | |||||
| info := &relaycommon.RelayInfo{ | |||||
| UsingGroup: "premium", | |||||
| UserGroup: "vip", | |||||
| } | |||||
| groupRatioInfo := HandleGroupRatio(c, info) | |||||
| assert.Equal(t, 0.8, groupRatioInfo.GroupRatio) | |||||
| assert.Equal(t, 0.8, groupRatioInfo.GroupSpecialRatio) | |||||
| assert.True(t, groupRatioInfo.HasSpecialRatio) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{}`)) | |||||
| } | |||||
| func TestHandleGroupRatio_AutoGroupOverridesUsingGroup(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"vip-channel":0.6}`)) | |||||
| c := buildTestContext(t) | |||||
| // Simulate: distributor set auto_group = "vip-channel" in context | |||||
| c.Set("auto_group", "vip-channel") | |||||
| info := &relaycommon.RelayInfo{ | |||||
| UsingGroup: "default", // will be overridden | |||||
| UserGroup: "default", | |||||
| } | |||||
| groupRatioInfo := HandleGroupRatio(c, info) | |||||
| // UsingGroup should be updated to "vip-channel" | |||||
| assert.Equal(t, "vip-channel", info.UsingGroup) | |||||
| // GroupRatio should reflect vip-channel's ratio | |||||
| assert.Equal(t, 0.6, groupRatioInfo.GroupRatio) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||||
| } | |||||
| func TestHandleGroupRatio_AutoGroupWithSpecialRatio(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"premium":1.0}`)) | |||||
| require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{"vip":{"premium":0.5}}`)) | |||||
| c := buildTestContext(t) | |||||
| c.Set("auto_group", "premium") | |||||
| info := &relaycommon.RelayInfo{ | |||||
| UsingGroup: "default", | |||||
| UserGroup: "vip", | |||||
| } | |||||
| groupRatioInfo := HandleGroupRatio(c, info) | |||||
| assert.Equal(t, "premium", info.UsingGroup) | |||||
| // Should use the special ratio (0.5) from vip->premium, not the generic group ratio (1.0) | |||||
| assert.Equal(t, 0.5, groupRatioInfo.GroupRatio) | |||||
| assert.True(t, groupRatioInfo.HasSpecialRatio) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||||
| require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{}`)) | |||||
| } | |||||
| // --- ModelPriceHelper with group ratio --- | |||||
| func TestModelPriceHelper_GroupRatioAffectsPreConsume(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"discount":0.5}`)) | |||||
| c := buildTestContext(t) | |||||
| info := &relaycommon.RelayInfo{ | |||||
| OriginModelName: testModelUsePriceSwitch, | |||||
| UsingGroup: "discount", | |||||
| UserGroup: "default", | |||||
| UserSetting: dto.UserSetting{}, | |||||
| } | |||||
| priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) | |||||
| require.NoError(t, err) | |||||
| // model price 0.5, group ratio 0.5 => preConsume = 0.5 * QuotaPerUnit * 0.5 | |||||
| assert.Equal(t, 0.5, priceData.GroupRatioInfo.GroupRatio) | |||||
| assert.True(t, priceData.UsePrice) | |||||
| // 0.5 (modelPrice) * 500000 (QuotaPerUnit) * 0.5 (groupRatio) = 125000 | |||||
| expectedQuota := int(0.5 * float64(common.QuotaPerUnit) * 0.5) | |||||
| assert.Equal(t, expectedQuota, priceData.QuotaToPreConsume) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||||
| } | |||||
| func TestModelPriceHelper_RatioModeWithGroupRatio(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| // Remove model price so it falls back to ratio mode | |||||
| require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"double":2.0}`)) | |||||
| c := buildTestContext(t) | |||||
| info := &relaycommon.RelayInfo{ | |||||
| OriginModelName: testModelUsePriceSwitch, | |||||
| UsingGroup: "double", | |||||
| UserGroup: "default", | |||||
| UserSetting: dto.UserSetting{}, | |||||
| } | |||||
| priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) | |||||
| require.NoError(t, err) | |||||
| assert.False(t, priceData.UsePrice) | |||||
| assert.Equal(t, 15.0, priceData.ModelRatio) | |||||
| assert.Equal(t, 2.0, priceData.GroupRatioInfo.GroupRatio) | |||||
| // ratio = modelRatio(15) * groupRatio(2.0) = 30 | |||||
| // preConsumedQuota = int(preConsumedTokens * 30) | |||||
| expectedRatio := 15.0 * 2.0 | |||||
| assert.Equal(t, expectedRatio, priceData.ModelRatio*priceData.GroupRatioInfo.GroupRatio) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||||
| } | |||||
| func TestModelPriceHelper_FreeGroup(t *testing.T) { | |||||
| setupUsePriceSwitchTest(t) | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"free":0}`)) | |||||
| // Disable pre-consume for free models so freeModel flag is set | |||||
| origSetting := operation_setting.GetQuotaSetting() | |||||
| operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = false | |||||
| t.Cleanup(func() { | |||||
| operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = origSetting.EnableFreeModelPreConsume | |||||
| }) | |||||
| c := buildTestContext(t) | |||||
| info := &relaycommon.RelayInfo{ | |||||
| OriginModelName: testModelUsePriceSwitch, | |||||
| UsingGroup: "free", | |||||
| UserGroup: "default", | |||||
| UserSetting: dto.UserSetting{}, | |||||
| } | |||||
| priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) | |||||
| require.NoError(t, err) | |||||
| // group ratio = 0 means free model | |||||
| assert.Equal(t, 0.0, priceData.GroupRatioInfo.GroupRatio) | |||||
| assert.True(t, priceData.FreeModel) | |||||
| assert.Equal(t, 0, priceData.QuotaToPreConsume) | |||||
| // cleanup | |||||
| require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) | |||||
| } | |||||
| @@ -0,0 +1,120 @@ | |||||
| package service | |||||
| import ( | |||||
| "net/http" | |||||
| "net/http/httptest" | |||||
| "testing" | |||||
| "github.com/QuantumNous/new-api/setting" | |||||
| "github.com/gin-gonic/gin" | |||||
| "github.com/stretchr/testify/assert" | |||||
| "github.com/stretchr/testify/require" | |||||
| ) | |||||
| func buildRetryContext(t *testing.T) *gin.Context { | |||||
| t.Helper() | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| req, err := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil) | |||||
| require.NoError(t, err) | |||||
| c.Request = req | |||||
| return c | |||||
| } | |||||
| // --- RetryParam tests --- | |||||
| func TestRetryParam_IncreaseRetry(t *testing.T) { | |||||
| p := &RetryParam{} | |||||
| assert.Equal(t, 0, p.GetRetry()) | |||||
| p.IncreaseRetry() | |||||
| assert.Equal(t, 1, p.GetRetry()) | |||||
| p.IncreaseRetry() | |||||
| assert.Equal(t, 2, p.GetRetry()) | |||||
| } | |||||
| func TestRetryParam_ResetRetryNextTry(t *testing.T) { | |||||
| p := &RetryParam{} | |||||
| p.IncreaseRetry() | |||||
| assert.Equal(t, 1, p.GetRetry()) | |||||
| // ResetRetryNextTry should cause next IncreaseRetry to be a no-op | |||||
| p.ResetRetryNextTry() | |||||
| p.IncreaseRetry() | |||||
| assert.Equal(t, 1, p.GetRetry()) // still 1 | |||||
| p.IncreaseRetry() | |||||
| assert.Equal(t, 2, p.GetRetry()) // now increases | |||||
| } | |||||
| func TestRetryParam_SetRetry(t *testing.T) { | |||||
| p := &RetryParam{} | |||||
| p.SetRetry(5) | |||||
| assert.Equal(t, 5, p.GetRetry()) | |||||
| } | |||||
| // --- Auto group selection (no DB, tests with no channels available) --- | |||||
| func TestCacheGetRandomSatisfiedChannel_AutoGroup_NoAutoGroups(t *testing.T) { | |||||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`)) | |||||
| t.Cleanup(func() { | |||||
| _ = setting.UpdateAutoGroupsByJsonString(`["default"]`) | |||||
| _ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`) | |||||
| }) | |||||
| c := buildRetryContext(t) | |||||
| p := &RetryParam{ | |||||
| Ctx: c, | |||||
| TokenGroup: "auto", | |||||
| ModelName: "gpt-4", | |||||
| } | |||||
| ch, group, err := CacheGetRandomSatisfiedChannel(p) | |||||
| assert.Nil(t, ch) | |||||
| assert.Error(t, err) | |||||
| assert.Equal(t, "auto", group) | |||||
| } | |||||
| // --- GetUserAutoGroup tests --- | |||||
| func TestGetUserAutoGroup_Basic(t *testing.T) { | |||||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip"]`)) | |||||
| require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)) | |||||
| t.Cleanup(func() { | |||||
| _ = setting.UpdateAutoGroupsByJsonString(`["default"]`) | |||||
| _ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`) | |||||
| }) | |||||
| // default user can use both groups | |||||
| autoGroups := GetUserAutoGroup("default") | |||||
| assert.Equal(t, []string{"default", "vip"}, autoGroups) | |||||
| } | |||||
| func TestGetUserAutoGroup_VipUserWithExtraGroups(t *testing.T) { | |||||
| require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip","premium"]`)) | |||||
| require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组","premium":"高级分组"}`)) | |||||
| t.Cleanup(func() { | |||||
| _ = setting.UpdateAutoGroupsByJsonString(`["default"]`) | |||||
| _ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`) | |||||
| }) | |||||
| autoGroups := GetUserAutoGroup("vip") | |||||
| assert.Equal(t, []string{"default", "vip", "premium"}, autoGroups) | |||||
| } | |||||
| // --- GetUserGroupRatio tests --- | |||||
| func TestGetUserGroupRatio_SpecialRatio(t *testing.T) { | |||||
| // When user has a special ratio for a group, it should be used | |||||
| // The default groupGroupRatio has vip->edit_this = 0.9 | |||||
| ratio := GetUserGroupRatio("vip", "edit_this") | |||||
| assert.Equal(t, 0.9, ratio) | |||||
| } | |||||
| func TestGetUserGroupRatio_FallbackToGroupRatio(t *testing.T) { | |||||
| // When no special ratio exists, fall back to group ratio | |||||
| ratio := GetUserGroupRatio("default", "default") | |||||
| assert.Equal(t, 1.0, ratio) // default group has ratio 1 | |||||
| } | |||||
| @@ -0,0 +1,126 @@ | |||||
| package ratio_setting | |||||
| import ( | |||||
| "testing" | |||||
| "github.com/stretchr/testify/assert" | |||||
| "github.com/stretchr/testify/require" | |||||
| ) | |||||
| const defaultGroupRatioJSON = `{"default":1,"vip":1,"svip":1}` | |||||
| const defaultGroupGroupRatioJSON = `{"vip":{"edit_this":0.9}}` | |||||
| // resetGroupRatios restores the default group ratios after each test. | |||||
| func resetGroupRatios(t *testing.T) { | |||||
| t.Helper() | |||||
| t.Cleanup(func() { | |||||
| _ = UpdateGroupRatioByJSONString(defaultGroupRatioJSON) | |||||
| _ = UpdateGroupGroupRatioByJSONString(defaultGroupGroupRatioJSON) | |||||
| }) | |||||
| } | |||||
| // --- GetGroupRatio --- | |||||
| func TestGetGroupRatio_DefaultGroups(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| assert.Equal(t, 1.0, GetGroupRatio("default")) | |||||
| assert.Equal(t, 1.0, GetGroupRatio("vip")) | |||||
| assert.Equal(t, 1.0, GetGroupRatio("svip")) | |||||
| } | |||||
| func TestGetGroupRatio_UnknownGroup(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| // 不存在的分组应返回 1(降级默认值) | |||||
| assert.Equal(t, 1.0, GetGroupRatio("nonexistent")) | |||||
| } | |||||
| func TestGetGroupRatio_UpdateAndRead(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| err := UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"test-group":0.8,"free":0}`) | |||||
| require.NoError(t, err) | |||||
| assert.Equal(t, 0.8, GetGroupRatio("test-group")) | |||||
| assert.Equal(t, 0.0, GetGroupRatio("free")) | |||||
| // 清除自定义分组后应回退到默认值 1 | |||||
| require.NoError(t, UpdateGroupRatioByJSONString(defaultGroupRatioJSON)) | |||||
| assert.Equal(t, 1.0, GetGroupRatio("test-group")) | |||||
| } | |||||
| func TestGetGroupRatio_OverrideDefault(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| require.NoError(t, UpdateGroupRatioByJSONString(`{"default":1,"vip":0.5,"svip":1}`)) | |||||
| assert.Equal(t, 0.5, GetGroupRatio("vip")) | |||||
| } | |||||
| func TestGetGroupRatio_InvalidJSON(t *testing.T) { | |||||
| err := UpdateGroupRatioByJSONString(`not json`) | |||||
| assert.Error(t, err) | |||||
| } | |||||
| func TestCheckGroupRatio_Negative(t *testing.T) { | |||||
| err := CheckGroupRatio(`{"bad":-0.5}`) | |||||
| assert.Error(t, err) | |||||
| } | |||||
| func TestCheckGroupRatio_Valid(t *testing.T) { | |||||
| err := CheckGroupRatio(`{"ok":0.5,"zero":0}`) | |||||
| assert.NoError(t, err) | |||||
| } | |||||
| // --- GetGroupGroupRatio --- | |||||
| func TestGetGroupGroupRatio_NoSpecialRatio(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| // default 用户使用 default 分组,没有特殊倍率 | |||||
| ratio, ok := GetGroupGroupRatio("default", "default") | |||||
| assert.False(t, ok) | |||||
| assert.Equal(t, -1.0, ratio) | |||||
| } | |||||
| func TestGetGroupGroupRatio_DefaultSpecialRatio(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| // 内置默认: vip 用户使用 edit_this 分组 => 0.9 | |||||
| ratio, ok := GetGroupGroupRatio("vip", "edit_this") | |||||
| assert.True(t, ok) | |||||
| assert.Equal(t, 0.9, ratio) | |||||
| } | |||||
| func TestGetGroupGroupRatio_UpdateAndRead(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| err := UpdateGroupGroupRatioByJSONString(`{"premium":{"groupA":0.7,"groupB":1.2}}`) | |||||
| require.NoError(t, err) | |||||
| ratio, ok := GetGroupGroupRatio("premium", "groupA") | |||||
| assert.True(t, ok) | |||||
| assert.Equal(t, 0.7, ratio) | |||||
| ratio, ok = GetGroupGroupRatio("premium", "groupB") | |||||
| assert.True(t, ok) | |||||
| assert.Equal(t, 1.2, ratio) | |||||
| // 不存在的组合 | |||||
| _, ok = GetGroupGroupRatio("premium", "nonexistent") | |||||
| assert.False(t, ok) | |||||
| _, ok = GetGroupGroupRatio("nonexistent", "groupA") | |||||
| assert.False(t, ok) | |||||
| } | |||||
| // --- GetGroupRatioCopy --- | |||||
| func TestGetGroupRatioCopy_ReturnsCurrentState(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| require.NoError(t, UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"copy-test":2.5}`)) | |||||
| copy := GetGroupRatioCopy() | |||||
| assert.Equal(t, 2.5, copy["copy-test"]) | |||||
| } | |||||
| // --- ContainsGroupRatio --- | |||||
| func TestContainsGroupRatio(t *testing.T) { | |||||
| resetGroupRatios(t) | |||||
| assert.True(t, ContainsGroupRatio("default")) | |||||
| assert.True(t, ContainsGroupRatio("vip")) | |||||
| assert.False(t, ContainsGroupRatio("nonexistent_group")) | |||||
| } | |||||
| @@ -181,39 +181,37 @@ const ModelPricingTable = ({ | |||||
| ); | ); | ||||
| }; | }; | ||||
| // 分组价格已隐藏 | |||||
| // return ( | |||||
| // <Card className='!rounded-2xl shadow-sm border-0'> | |||||
| // <div className='flex items-center mb-4'> | |||||
| // <Avatar size='small' color='orange' className='mr-2 shadow-md'> | |||||
| // <IconCoinMoneyStroked size={16} /> | |||||
| // </Avatar> | |||||
| // <div> | |||||
| // <Text className='text-lg font-medium'>{t('分组价格')}</Text> | |||||
| // <div className='text-xs text-gray-600'> | |||||
| // {t('不同用户分组的价格信息')} | |||||
| // </div> | |||||
| // </div> | |||||
| // </div> | |||||
| // {autoChain.length > 0 && ( | |||||
| // <div className='flex flex-wrap items-center gap-1 mb-4'> | |||||
| // <span className='text-sm text-gray-600'>{t('auto分组调用链路')}</span> | |||||
| // <span className='text-sm'>→</span> | |||||
| // {autoChain.map((g, idx) => ( | |||||
| // <React.Fragment key={g}> | |||||
| // <Tag color='white' size='small' shape='circle'> | |||||
| // {g} | |||||
| // {t('分组')} | |||||
| // </Tag> | |||||
| // {idx < autoChain.length - 1 && <span className='text-sm'>→</span>} | |||||
| // </React.Fragment> | |||||
| // ))} | |||||
| // </div> | |||||
| // )} | |||||
| // {renderGroupPriceTable()} | |||||
| // </Card> | |||||
| // ); | |||||
| return null; | |||||
| return ( | |||||
| <Card className='!rounded-2xl shadow-sm border-0'> | |||||
| <div className='flex items-center mb-4'> | |||||
| <Avatar size='small' color='orange' className='mr-2 shadow-md'> | |||||
| <IconCoinMoneyStroked size={16} /> | |||||
| </Avatar> | |||||
| <div> | |||||
| <Text className='text-lg font-medium'>{t('分组价格')}</Text> | |||||
| <div className='text-xs text-gray-600'> | |||||
| {t('不同用户分组的价格信息')} | |||||
| </div> | |||||
| </div> | |||||
| </div> | |||||
| {autoChain.length > 0 && ( | |||||
| <div className='flex flex-wrap items-center gap-1 mb-4'> | |||||
| <span className='text-sm text-gray-600'>{t('auto分组调用链路')}</span> | |||||
| <span className='text-sm'>→</span> | |||||
| {autoChain.map((g, idx) => ( | |||||
| <React.Fragment key={g}> | |||||
| <Tag color='white' size='small' shape='circle'> | |||||
| {g} | |||||
| {t('分组')} | |||||
| </Tag> | |||||
| {idx < autoChain.length - 1 && <span className='text-sm'>→</span>} | |||||
| </React.Fragment> | |||||
| ))} | |||||
| </div> | |||||
| )} | |||||
| {renderGroupPriceTable()} | |||||
| </Card> | |||||
| ); | |||||
| }; | }; | ||||
| export default ModelPricingTable; | export default ModelPricingTable; | ||||