235 lines
12 KiB
Go
235 lines
12 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
"net/http/httptest"
|
|
)
|
|
|
|
func TestGroupDegradationDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
|
t.Skip("TEST_MYSQL_DSN unset")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "postgres":
|
|
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
|
t.Skip("TEST_POSTGRES_DSN unset")
|
|
}
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("degradation_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
oldDB := model.DB
|
|
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
|
model.DB = db
|
|
common.SetDatabaseTypes(common.DatabaseType(dialect), oldLogType)
|
|
rows := []any{&model.Option{}, &model.Channel{}, &model.SystemTask{}, &model.DetectionSchedule{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.Token{}}
|
|
require.NoError(t, db.AutoMigrate(rows...))
|
|
t.Cleanup(func() {
|
|
for _, row := range rows {
|
|
require.NoError(t, db.Migrator().DropTable(row))
|
|
}
|
|
model.DB = oldDB
|
|
common.SetDatabaseTypes(oldType, oldLogType)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
groupDegradationSnapshot.latest = nil
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
require.NoError(t, conn.Close())
|
|
})
|
|
var version string
|
|
query := "select version()"
|
|
if dialect == "sqlite" {
|
|
query = "select sqlite_version()"
|
|
}
|
|
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
|
t.Log("database version:", version)
|
|
now := time.Now().Unix()
|
|
require.NoError(t, db.Create(&model.Channel{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1, Key: "test-only"}).Error)
|
|
body, err := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ID: "mismatch", ChannelID: 1, StartedAt: now - 30, FinishedAt: now - 10, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 2}}}})
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.Create(&model.SystemTask{TaskID: "degradation-test", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 30, State: string(body)}).Error)
|
|
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}}
|
|
encoded, err := common.Marshal(policy)
|
|
require.NoError(t, err)
|
|
token := model.Token{UserId: 1, Key: "degradation-test-key", RoutingPolicy: string(encoded)}
|
|
require.NoError(t, db.Create(&token).Error)
|
|
require.NoError(t, db.First(&token, token.Id).Error)
|
|
stored, err := token.GetRoutingPolicy()
|
|
require.NoError(t, err)
|
|
require.True(t, stored.SkipDegradedGroups)
|
|
ctx := func() *gin.Context {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Set(model.TokenRoutingContextKey, stored)
|
|
return c
|
|
}
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
smart := ctx()
|
|
smartPolicy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart"}
|
|
smart.Set(model.TokenRoutingContextKey, smartPolicy)
|
|
assert.True(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
|
|
smartPolicy.AllowDegraded = true
|
|
assert.False(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
|
|
smartPolicy.AllowDegraded = false
|
|
assert.False(t, requestGroupDegraded(smart, "unmeasured-model", "default"))
|
|
manualResult, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2})
|
|
require.NoError(t, err)
|
|
manual := model.DetectionRun{ID: "manual-recovery", RequestKey: "manual-recovery", ChannelID: 1, Kind: "fingerprint", UserID: 1, CreatedAt: now, Settled: true, Status: "succeeded", Result: string(manualResult)}
|
|
require.NoError(t, db.Create(&manual).Error)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
pelican := model.DetectionRun{ID: "pelican-degraded", RequestKey: "pelican-degraded", ChannelID: 1, Kind: "pelican", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
|
|
require.NoError(t, db.Create(&pelican).Error)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
require.NoError(t, db.Model(&pelican).Update("quality", "normal").Error)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
require.NoError(t, db.Delete(&manual).Error)
|
|
require.NoError(t, model.SaveStatusPageGroup(ctx().Request.Context(), "default", nil))
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGroupDegradationUsesLatestValidBeforeHistoryTruncation(t *testing.T) {
|
|
now := time.Now()
|
|
channels := []model.Channel{{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1}}
|
|
rounds := []AstraSamplingRound{{ID: "valid", ChannelID: 1, StartedAt: now.Unix() - 100, FinishedAt: now.Unix() - 90, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 1}}}
|
|
for i := range 15 {
|
|
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix() - int64(i), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "error"}})
|
|
}
|
|
body, err := common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
require.NoError(t, err)
|
|
tasks := []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(body)}}
|
|
summary := AggregateAstraSampling(channels, tasks, now)["default"]
|
|
require.NotNil(t, summary.LatestValid)
|
|
assert.Equal(t, "gpt-5.5", summary.LatestValid.Closest)
|
|
assert.Len(t, summary.History, 12)
|
|
rounds = append(rounds, AstraSamplingRound{ID: "recovered", ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2}})
|
|
body, err = common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
require.NoError(t, err)
|
|
tasks[0].State = string(body)
|
|
summary = AggregateAstraSampling(channels, tasks, now)["default"]
|
|
require.NotNil(t, summary.LatestValid)
|
|
assert.Equal(t, fingerprint.Model, summary.LatestValid.Closest)
|
|
reset := aggregateAstraSampling(channels, tasks, now, map[string]model.StatusPageGroup{"default": {ResetBefore: now.Unix()}})
|
|
assert.Nil(t, reset["default"].LatestValid)
|
|
}
|
|
|
|
func TestGroupDegradationOverridesExistingPoolAffinity(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6401, "default", fingerprint.Model)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6402, "vip", fingerprint.Model)
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
previousRedis := common.RedisEnabled
|
|
common.RedisEnabled = false
|
|
t.Cleanup(func() { common.RedisEnabled = previousRedis; _ = poolBindings().Purge() })
|
|
require.NoError(t, poolBindings().Purge())
|
|
config := operation_setting.GetChannelAffinitySetting()
|
|
previous := *config
|
|
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "degradation-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
|
|
t.Cleanup(func() { *config = previous })
|
|
ctx := func(degraded bool) *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Request.Header.Set("Session-Id", "degradation-session")
|
|
c.Set("id", 100)
|
|
c.Set("token_id", 100)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
|
|
rounds := map[string]AstraSamplingHistory{}
|
|
if degraded {
|
|
rounds["default"] = AstraSamplingHistory{At: time.Now().Unix(), Closest: "gpt-5.5", ValidSamples: 1}
|
|
}
|
|
c.Set("pool_group_fingerprints", rounds)
|
|
return c
|
|
}
|
|
first := ctx(false)
|
|
selected, group, err := SelectChannelForRequest(first, fingerprint.Model, &RetryParam{Ctx: first, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.Nil(t, err)
|
|
require.NotNil(t, selected)
|
|
require.Equal(t, 6401, selected.Id)
|
|
first.Set("channel_id", selected.Id)
|
|
common.SetContextKey(first, constant.ContextKeyAutoGroup, group)
|
|
RequestPolicy(first).Successful = true
|
|
RecordPoolAffinity(first)
|
|
bound, _ := PreferredPoolChannel(ctx(false), fingerprint.Model)
|
|
require.NotNil(t, bound)
|
|
next := ctx(true)
|
|
selected, group, err = SelectChannelForRequest(next, fingerprint.Model, &RetryParam{Ctx: next, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.Nil(t, err)
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, 6402, selected.Id)
|
|
assert.Equal(t, "vip", group)
|
|
}
|
|
|
|
func TestGroupDegradationSkipsPoolGroupsAndRespectsUnknownRecoveryAndOptOut(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", fingerprint.Model)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", fingerprint.Model)
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
for _, test := range []struct {
|
|
name string
|
|
enabled bool
|
|
closest string
|
|
age time.Duration
|
|
degradeBackup bool
|
|
want int
|
|
}{
|
|
{"latest mismatch switches", true, "gpt-5.5", time.Minute, false, 6302},
|
|
{"recovered group returns", true, fingerprint.Model, time.Minute, false, 6301},
|
|
{"unknown stays eligible", true, "", time.Minute, false, 6301},
|
|
{"valid result persists beyond window", true, "gpt-5.5", 7 * time.Hour, false, 6302},
|
|
{"disabled keeps routing", false, "gpt-5.5", time.Minute, false, 6301},
|
|
{"all degraded has no route", true, "gpt-5.5", time.Minute, true, 0},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
c := setupGroupPoolContext(t)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: test.enabled, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
|
|
rounds := map[string]AstraSamplingHistory{"default": {At: time.Now().Add(-test.age).Unix(), Closest: test.closest, ValidSamples: 1}}
|
|
if test.degradeBackup {
|
|
rounds["vip"] = rounds["default"]
|
|
}
|
|
c.Set("pool_group_fingerprints", rounds)
|
|
selected, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.NoError(t, err)
|
|
if test.want == 0 {
|
|
assert.Nil(t, selected)
|
|
return
|
|
}
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, test.want, selected.Id)
|
|
})
|
|
}
|
|
}
|