cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
70 lines
3.1 KiB
Go
70 lines
3.1 KiB
Go
package service
|
|
|
|
import (
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestTokenRoutingPoolPreferenceExclusionAndAffinityIsolation(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
const modelName = "pool-model"
|
|
createChannelSelectAutoGroupsChannel(t, db, 6101, "default", modelName)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6102, "vip", modelName)
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
config := operation_setting.GetChannelAffinitySetting()
|
|
previous := *config
|
|
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "pool-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
|
|
t.Cleanup(func() { *config = previous; _ = poolBindings().Purge() })
|
|
ctx := func(token int, session string) *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Request.Header.Set("Session-Id", session)
|
|
c.Set("id", 100)
|
|
c.Set("token_id", token)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}, ModelRules: []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}})
|
|
return c
|
|
}
|
|
c := ctx(100, "session-a")
|
|
param := &RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)}
|
|
first, group, err := SelectChannelForRequest(c, modelName, param)
|
|
require.Nil(t, err)
|
|
require.NotNil(t, first)
|
|
require.Equal(t, 6102, first.Id)
|
|
require.Equal(t, "vip", group)
|
|
AppendUsedChannel(c, first.Id)
|
|
second, group, err2 := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err2)
|
|
require.NotNil(t, second)
|
|
require.Equal(t, 6101, second.Id)
|
|
c.Set("channel_id", second.Id)
|
|
common.SetContextKey(c, constant.ContextKeyAutoGroup, group)
|
|
RequestPolicy(c).Successful = true
|
|
RecordPoolAffinity(c)
|
|
next := ctx(100, "session-a")
|
|
bound, boundGroup := PreferredPoolChannel(next, modelName)
|
|
require.NotNil(t, bound)
|
|
require.Equal(t, 6101, bound.Id)
|
|
require.Equal(t, "default", boundGroup)
|
|
isolated, _ := PreferredPoolChannel(ctx(101, "session-a"), modelName)
|
|
require.Nil(t, isolated)
|
|
isolated, _ = PreferredPoolChannel(ctx(100, "session-b"), modelName)
|
|
require.Nil(t, isolated)
|
|
isolated, _ = PreferredPoolChannel(ctx(100, "session-a"), "another-model")
|
|
require.Nil(t, isolated)
|
|
ClearPoolAffinity(next)
|
|
bound, _ = PreferredPoolChannel(ctx(100, "session-a"), modelName)
|
|
require.Nil(t, bound)
|
|
AppendUsedChannel(c, second.Id)
|
|
last, _, err2 := CacheGetRandomSatisfiedChannel(param)
|
|
require.NoError(t, err2)
|
|
require.Nil(t, last)
|
|
}
|