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/model" | |||
| 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/types" | |||
| "github.com/gin-gonic/gin" | |||
| @@ -148,3 +149,194 @@ func TestUpdatePriceDataForChannelPricingOnlyAppliesUserChannelRatio(t *testing. | |||
| assert.Equal(t, 800, info.PriceData.QuotaToPreConsume) | |||
| 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; | |||