feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.

service/group_pool_test.go

@@ -0,0 +1,153 @@
+package service
+
+import (
+	"net/http/httptest"
+	"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/require"
+)
+
+func setupGroupPoolContext(t *testing.T) *gin.Context {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+	ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
+	common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
+	common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "default")
+	common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true)
+	common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
+	common.SetContextKey(ctx, constant.ContextKeyTokenGroupPool, []model.TokenGroupEntry{{Group: "vip", Priority: 1}})
+	return ctx
+}
+
+func TestBuildGroupPoolTiersOrdersByPriorityAndPrependsPrimary(t *testing.T) {
+	setupChannelSelectAutoGroupsTest(t)
+	ctx := setupGroupPoolContext(t)
+
+	tiers := buildGroupPoolTiers(ctx, "default", "default")
+	require.Len(t, tiers, 2)
+	require.Equal(t, []string{"default"}, tiers[0].groups)
+	require.Equal(t, []string{"vip"}, tiers[1].groups)
+	require.Less(t, tiers[0].priority, tiers[1].priority)
+}
+
+func TestBuildGroupPoolTiersGroupsEqualPriorityTogether(t *testing.T) {
+	setupChannelSelectAutoGroupsTest(t)
+	ctx := setupGroupPoolContext(t)
+	common.SetContextKey(ctx, constant.ContextKeyTokenGroupPool, []model.TokenGroupEntry{
+		{Group: "default", Priority: 1},
+		{Group: "vip", Priority: 1},
+	})
+
+	tiers := buildGroupPoolTiers(ctx, "default", "default")
+	require.Len(t, tiers, 1)
+	require.Equal(t, []string{"default", "vip"}, tiers[0].groups)
+}
+
+func TestNextTierIndexAlternatesWithinTier(t *testing.T) {
+	key := "round-robin-" + t.Name()
+	require.Equal(t, 0, nextTierIndex(key, 2))
+	require.Equal(t, 1, nextTierIndex(key, 2))
+	require.Equal(t, 0, nextTierIndex(key, 2))
+	require.Equal(t, 1, nextTierIndex(key, 2))
+	require.Equal(t, 0, nextTierIndex(key, 1))
+}
+
+func TestGroupPoolSwitchesGroupAfterRequestFailure(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	const modelName = "group-pool-switch-after-failure"
+	createChannelSelectAutoGroupsChannel(t, db, 3101, "default", modelName)
+	createChannelSelectAutoGroupsChannel(t, db, 3102, "vip", modelName)
+	model.InitChannelCache()
+	ResetGroupAvailability()
+	t.Cleanup(ResetGroupAvailability)
+
+	ctx := setupGroupPoolContext(t)
+	param := &RetryParam{Ctx: ctx, TokenGroup: "default", ModelName: modelName}
+
+	first, firstGroup, err := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err)
+	require.NotNil(t, first)
+	require.Equal(t, 3101, first.Id)
+	require.Equal(t, "default", firstGroup)
+
+	MarkRequestGroupFailed(ctx, firstGroup)
+
+	second, secondGroup, err := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err)
+	require.NotNil(t, second)
+	require.Equal(t, 3102, second.Id)
+	require.Equal(t, "vip", secondGroup)
+}
+
+func TestGroupPoolAdvancesTierOnRetry(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	const modelName = "group-pool-advance-on-retry"
+	createChannelSelectAutoGroupsChannel(t, db, 3301, "default", modelName)
+	createChannelSelectAutoGroupsChannel(t, db, 3302, "vip", modelName)
+	model.InitChannelCache()
+	ResetGroupAvailability()
+	t.Cleanup(ResetGroupAvailability)
+
+	ctx := setupGroupPoolContext(t)
+	param := &RetryParam{Ctx: ctx, TokenGroup: "default", ModelName: modelName}
+
+	first, firstGroup, err := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err)
+	require.NotNil(t, first)
+	require.Equal(t, 3301, first.Id)
+	require.Equal(t, "default", firstGroup)
+
+	param.IncreaseRetry()
+
+	second, secondGroup, err := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err)
+	require.NotNil(t, second)
+	require.Equal(t, 3302, second.Id)
+	require.Equal(t, "vip", secondGroup)
+}
+
+func TestGroupPoolSkipsUnavailableGroup(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	const modelName = "group-pool-skip-unavailable"
+	createChannelSelectAutoGroupsChannel(t, db, 3201, "default", modelName)
+	createChannelSelectAutoGroupsChannel(t, db, 3202, "vip", modelName)
+	model.InitChannelCache()
+	ResetGroupAvailability()
+	t.Cleanup(ResetGroupAvailability)
+
+	ctx := setupGroupPoolContext(t)
+	for i := 0; i < groupAvailabilityFailureBurst; i++ {
+		RecordGroupOutcome("default", false)
+	}
+	require.False(t, IsGroupAvailable("default"))
+
+	param := &RetryParam{Ctx: ctx, TokenGroup: "default", ModelName: modelName}
+	channel, group, err := CacheGetRandomSatisfiedChannel(param)
+	require.NoError(t, err)
+	require.NotNil(t, channel)
+	require.Equal(t, 3202, channel.Id)
+	require.Equal(t, "vip", group)
+}
+
+func TestGroupAvailabilityRecoversAfterCooldown(t *testing.T) {
+	ResetGroupAvailability()
+	t.Cleanup(ResetGroupAvailability)
+
+	for i := 0; i < groupAvailabilityFailureBurst; i++ {
+		RecordGroupOutcome("recovery-group", false)
+	}
+	require.False(t, IsGroupAvailable("recovery-group"))
+
+	state := groupAvailabilityStateFor("recovery-group")
+	state.mu.Lock()
+	for i := range state.outcomes {
+		state.outcomes[i].at = state.outcomes[i].at.Add(-groupAvailabilityCooldown)
+	}
+	state.mu.Unlock()
+
+	require.True(t, IsGroupAvailable("recovery-group"))
+}