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

service/group_quality_test.go

@@ -0,0 +1,100 @@
+package service
+
+import (
+	"context"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/model"
+	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQualitySwitchesAndFailedReadCache(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	require.NoError(t, db.AutoMigrate(&model.Option{}, &model.SystemTask{}, &model.DetectionRun{}))
+	createChannelSelectAutoGroupsChannel(t, db, 6501, "default", fingerprint.Model)
+	old := time.Now().Add(-8 * time.Hour).Unix()
+	result, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", ValidSamples: 3, Closest: "gpt-6-sol"})
+	require.NoError(t, err)
+	run := model.DetectionRun{Result: string(result), ID: "old-valid", RequestKey: "old-valid", ChannelID: 6501, Kind: "fingerprint", Status: "succeeded", Settled: true, Quality: "degraded", CreatedAt: old}
+	require.NoError(t, db.Create(&run).Error)
+	require.NoError(t, db.Create(&model.DetectionRun{ID: "failed", RequestKey: "failed", ChannelID: 6501, Kind: "pelican", Status: "failed", Settled: true, CreatedAt: time.Now().Unix()}).Error)
+	groupDegradationSnapshot.latest = nil
+	InvalidateGroupQualityCache()
+	t.Cleanup(func() {
+		groupDegradationSnapshot.latest = nil
+		groupDegradationSnapshot.benchmark = nil
+		InvalidateGroupQualityCache()
+	})
+	ctx := func(allow, child bool) *gin.Context {
+		c, _ := gin.CreateTestContext(httptest.NewRecorder())
+		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AllowDegraded: allow, SkipDegradedNonAstra: child})
+		return c
+	}
+	for _, allow := range []bool{false, true} {
+		for _, child := range []bool{false, true} {
+			require.Equal(t, !allow, requestGroupDegraded(ctx(allow, child), fingerprint.Model, "default"))
+			require.Equal(t, !allow && child, requestGroupDegraded(ctx(allow, child), "gpt-6-sol", "default"))
+		}
+	}
+	recent, err := RecentDegradedGroups(context.Background())
+	require.NoError(t, err)
+	require.False(t, recent["default"])
+	require.NoError(t, db.Migrator().DropTable(&model.DetectionRun{}))
+	InvalidateGroupQualityCache()
+	require.True(t, requestGroupDegraded(ctx(false, false), fingerprint.Model, "default"), "read failure retains valid conclusion")
+	groupDegradationSnapshot.latest = nil
+	InvalidateGroupQualityCache()
+	cold := ctx(false, false)
+	require.True(t, requestGroupDegraded(cold, fingerprint.Model, "default"))
+	require.True(t, cold.GetBool("pool_degradation_unavailable"))
+	require.NoError(t, db.AutoMigrate(&model.DetectionRun{}))
+	InvalidateGroupQualityCache()
+	require.False(t, requestGroupDegraded(ctx(false, false), fingerprint.Model, "default"), "successful empty history remains eligible")
+}
+
+func TestAbandonedPoolProbeIsReleasedWithoutClearingFailures(t *testing.T) {
+	db := setupChannelSelectAutoGroupsTest(t)
+	createChannelSelectAutoGroupsChannel(t, db, 7001, "vip", "m")
+	model.InitChannelCache()
+	c := setupGroupPoolContext(t)
+	c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+	c.Set("id", 701)
+	c.Set("token_id", 701)
+	key := poolCooldownKey(c, "m", "default")
+	require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
+	t.Cleanup(func() { _, _ = poolCooldowns().DeleteMany([]string{key}) })
+	require.True(t, claimPoolGroup(c, "m", "default"))
+	require.True(t, poolGroupSkipped(c, "m", "default"))
+	ReleasePoolProbes(c)
+	require.False(t, poolGroupSkipped(c, "m", "default"))
+	require.True(t, claimPoolGroup(c, "m", "default"))
+	c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
+	RecordPoolRetryableFailure(c, "m", "default")
+	ReleasePoolProbes(c)
+	value, found, err := poolCooldowns().Get(key)
+	require.NoError(t, err)
+	require.True(t, found)
+	require.Greater(t, value.Until, time.Now().UnixMilli())
+}
+
+func TestManualPoolNonAstraDegradationSwitch(t *testing.T) {
+	for _, mode := range []string{"inherit", "custom"} {
+		for _, enabled := range []bool{false, true} {
+			for _, nonAstra := range []bool{false, true} {
+				c, _ := gin.CreateTestContext(httptest.NewRecorder())
+				c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, SkipDegradedGroups: enabled, SkipDegradedNonAstra: nonAstra})
+				c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{"default": {Status: "mismatch", ValidSamples: 3, Closest: "gpt-6-sol"}})
+				c.Set("pool_group_pelicans", map[string]DetectionHistory{"default": {Quality: "degraded"}})
+				require.Equal(t, enabled, requestGroupDegraded(c, fingerprint.Model, "default"))
+				require.Equal(t, enabled && nonAstra, requestGroupDegraded(c, "gpt-6-sol", "default"))
+				require.False(t, requestGroupDegraded(c, "gpt-6-sol", "untested"))
+			}
+		}
+	}
+}