new-api-custom/service/group_pool_test.go

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"))
}