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

service/pool_no_backup_test.go

@@ -0,0 +1,95 @@
+package service
+
+import (
+	"context"
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+	"net/http/httptest"
+	"testing"
+	"time"
+)
+
+func TestPoolWithoutAvailableBackupDoesNotCoolDown(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	createChannelSelectAutoGroupsChannel(t, db, 7801, "default", "m")
+	createChannelSelectAutoGroupsChannel(t, db, 7802, "vip", "other")
+	model.InitChannelCache()
+	ResetGroupAvailability()
+	t.Cleanup(func() { ResetGroupAvailability(); _ = poolCooldowns().Purge() })
+	makeCtx := func() *gin.Context {
+		c := setupGroupPoolContext(t)
+		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+		c.Set("token_id", 7801)
+		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
+		return c
+	}
+	c := makeCtx()
+	RecordPoolRetryableFailure(c, "m", "default")
+	RecordPoolRetryableFailure(c, "m", "default")
+	_, found, err := poolCooldowns().Get(poolCooldownKey(c, "m", "default"))
+	require.NoError(t, err)
+	require.False(t, found)
+	// An old cooldown and global health state must not strand a single eligible group.
+	require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(c, "m", "default"), poolCooldown{Until: time.Now().Add(10 * time.Minute).UnixMilli()}, 0))
+	for i := 0; i < 4; i++ {
+		RecordGroupOutcome("default", false)
+	}
+	next := makeCtx()
+	ch, group, err := selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(0)})
+	require.NoError(t, err)
+	require.NotNil(t, ch)
+	require.Equal(t, "default", group)
+	AppendUsedChannel(next, ch.Id)
+	RecordPoolRetryableFailure(next, "m", "default")
+	ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(1)})
+	require.NoError(t, err)
+	require.NotNil(t, ch)
+	AppendUsedChannel(next, ch.Id)
+	RecordPoolRetryableFailure(next, "m", "default")
+	ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(2)})
+	require.NoError(t, err)
+	require.Nil(t, ch, "do not loop indefinitely inside one request")
+}
+func TestAutomaticSpeedUsesLatestExactModelStreamingTest(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	require.NoError(t, db.AutoMigrate(&model.Log{}))
+	previous := model.LOG_DB
+	model.LOG_DB = db
+	t.Cleanup(func() { model.LOG_DB = previous })
+	now := time.Now().Unix()
+	samples := []model.Log{
+		{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 20, Other: `{"frt":500}`},
+		{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 10, Other: `{"frt":60000}`},
+		{ChannelId: 1, Type: 2, TokenId: 9, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now, Other: `{"frt":1}`},
+		{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 5, Other: `{"frt":1200}`},
+		{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "other", IsStream: true, CreatedAt: now, Other: `{"frt":10}`},
+		{ChannelId: 3, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: false, CreatedAt: now, Other: `{"frt":10}`},
+		{ChannelId: 4, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 2*3600, Other: `{"frt":10}`},
+	}
+	require.NoError(t, db.Create(&samples).Error)
+	speeds := latestModelTestSpeeds(context.Background(), "m")
+	require.Equal(t, map[int]float64{1: 60000, 2: 1200}, speeds)
+}
+
+func TestProbeAndStreamingCooldownsAreIsolated(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	createChannelSelectAutoGroupsChannel(t, db, 7901, "default", "m")
+	createChannelSelectAutoGroupsChannel(t, db, 7902, "vip", "m")
+	model.InitChannelCache()
+	t.Cleanup(func() { _ = poolCooldowns().Purge() })
+	makeCtx := func(stream bool) *gin.Context {
+		c := setupGroupPoolContext(t)
+		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+		c.Set("token_id", 7901)
+		c.Set("pool_request_stream", stream)
+		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
+		return c
+	}
+	probe := makeCtx(false)
+	RecordPoolRetryableFailure(probe, "m", "default")
+	RecordPoolRetryableFailure(probe, "m", "default")
+	require.True(t, poolGroupSkipped(makeCtx(false), "m", "default"))
+	require.False(t, poolGroupSkipped(makeCtx(true), "m", "default"))
+}