101 lines
4.8 KiB
Go
101 lines
4.8 KiB
Go
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"))
|
|
}
|
|
}
|
|
}
|
|
}
|