157 lines
7.2 KiB
Go
157 lines
7.2 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"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/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func adaptiveTestDB(t *testing.T) {
|
|
t.Helper()
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
conn.SetMaxOpenConns(1)
|
|
old := model.DB
|
|
model.DB = db
|
|
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
|
|
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.SystemTask{}, &model.DetectionSchedule{}, &model.DetectionRun{}))
|
|
}
|
|
|
|
func seedAutomaticFingerprint(t *testing.T, id string, at int64, closest string, force bool, running bool) {
|
|
t.Helper()
|
|
result := fingerprint.Result{Engine: fingerprint.Version, Status: "error"}
|
|
if closest != "" {
|
|
result = fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: closest, ValidSamples: 3}
|
|
if closest == fingerprint.Model {
|
|
result.Status = "match"
|
|
}
|
|
}
|
|
finished, status := at, model.SystemTaskStatusSucceeded
|
|
if running {
|
|
result.Status = "running"
|
|
finished = 0
|
|
status = model.SystemTaskStatusRunning
|
|
}
|
|
state, err := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ID: id, ChannelID: 1, StartedAt: at - 10, FinishedAt: finished, Result: result}}})
|
|
require.NoError(t, err)
|
|
payload, err := common.Marshal(map[string]bool{"force": force})
|
|
require.NoError(t, err)
|
|
require.NoError(t, model.DB.Create(&model.SystemTask{TaskID: id, Type: AstraSamplingTaskType, Status: status, State: string(state), Payload: string(payload), CreatedAt: at - 10, UpdatedAt: at}).Error)
|
|
}
|
|
|
|
func TestAdaptiveFingerprintBootstrapRecoveryAndManualIsolation(t *testing.T) {
|
|
adaptiveTestDB(t)
|
|
ctx := context.Background()
|
|
channel := model.Channel{Id: 1, Status: common.ChannelStatusEnabled, Group: "one,shared", Models: fingerprint.Model}
|
|
require.NoError(t, model.DB.Create(&channel).Error)
|
|
seedAutomaticFingerprint(t, "normal", 1000, fingerprint.Model, false, false)
|
|
seedAutomaticFingerprint(t, "unknown", 2000, "", false, false)
|
|
seedAutomaticFingerprint(t, "forced", 2100, "gpt-6-sol", true, false)
|
|
due, states, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(2899, 0))
|
|
require.NoError(t, err)
|
|
assert.Empty(t, due)
|
|
assert.Equal(t, "normal", states[1].LastValid)
|
|
assert.Equal(t, 15, states[1].IntervalMinutes)
|
|
due, _, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(2900, 0))
|
|
require.NoError(t, err)
|
|
require.Len(t, due, 1)
|
|
seedAutomaticFingerprint(t, "recovered", 3000, fingerprint.Model, false, false)
|
|
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(3000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 30, states[1].IntervalMinutes)
|
|
seedAutomaticFingerprint(t, "downgraded", 5000, "gpt-6-sol", false, false)
|
|
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(5000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 15, states[1].IntervalMinutes)
|
|
seedAutomaticFingerprint(t, "other-non-astra", 6000, "gpt-5.5", false, false)
|
|
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(6000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 30, states[1].IntervalMinutes)
|
|
// A completed result with a missing clock update is replayed once, not twice.
|
|
_, again, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(6000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, states, again)
|
|
seedAutomaticFingerprint(t, "abandoned", 7000, "", false, true)
|
|
due, _, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(10000, 0))
|
|
require.NoError(t, err)
|
|
assert.Empty(t, due)
|
|
require.NoError(t, model.DB.Model(&model.SystemTask{}).Where("task_id = ?", "abandoned").Updates(map[string]any{"status": model.SystemTaskStatusFailed, "updated_at": 10000}).Error)
|
|
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(10000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(10900), states[1].NextDueAt)
|
|
assert.Equal(t, "degraded", states[1].LastValid)
|
|
require.NoError(t, model.DB.Where("type = ?", AstraSamplingTaskType).Delete(&model.SystemTask{}).Error)
|
|
_, again, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(10000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, states, again)
|
|
}
|
|
|
|
func TestAdaptiveBenchmarkBootstrapInFlightAndRanges(t *testing.T) {
|
|
adaptiveTestDB(t)
|
|
ctx := context.Background()
|
|
channels := []model.Channel{{Id: 1, Group: "one,shared", Models: fingerprint.Model, Status: common.ChannelStatusEnabled}, {Id: 2, Group: "shared", Models: fingerprint.Model, Status: common.ChannelStatusEnabled}}
|
|
require.NoError(t, model.DB.Create(&channels).Error)
|
|
runs := []model.DetectionRun{
|
|
{ID: "auto1", RequestKey: "a", ChannelID: 1, Kind: "benchmark", Status: "succeeded", Quality: "normal", Settled: true, FinishedAt: 1000},
|
|
{ID: "auto2", RequestKey: "b", ChannelID: 2, Kind: "benchmark", Status: "failed", Settled: true, FinishedAt: 1000},
|
|
{ID: "paid", RequestKey: "c", UserID: 1, ChannelID: 1, Kind: "benchmark", Status: "succeeded", Quality: "degraded", Settled: true, FinishedAt: 1100},
|
|
{ID: "active", RequestKey: "d", ChannelID: 2, Kind: "benchmark", Status: "running"},
|
|
}
|
|
require.NoError(t, model.DB.Create(&runs).Error)
|
|
due, states, err := AutomaticDetectionChannelsDue(ctx, "benchmark", time.Unix(3000, 0))
|
|
require.NoError(t, err)
|
|
require.Len(t, due, 1)
|
|
assert.Equal(t, 1, due[0].Id)
|
|
assert.Equal(t, "normal", states[1].LastValid)
|
|
assert.Equal(t, 30, states[1].IntervalMinutes)
|
|
assert.Equal(t, 15, states[2].IntervalMinutes)
|
|
ranges, err := DetectionScheduleRanges(ctx, channels, "benchmark")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, [2]int{15, 30}, ranges["shared"])
|
|
assert.Equal(t, [2]int{30, 30}, ranges["one"])
|
|
_, fp, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(3000, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 30, fp[2].IntervalMinutes)
|
|
assert.Equal(t, "", fp[1].LastValid)
|
|
require.NoError(t, model.DB.Model(&model.DetectionRun{}).Where("id = ?", "active").Updates(map[string]any{"status": "succeeded", "quality": "degraded", "settled": true, "finished_at": 3000}).Error)
|
|
due, states, err = AutomaticDetectionChannelsDue(ctx, "benchmark", time.Unix(4799, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 30, states[2].IntervalMinutes)
|
|
require.Len(t, due, 1)
|
|
due, _, err = AutomaticDetectionChannelsDue(ctx, "benchmark", time.Unix(4800, 0))
|
|
require.NoError(t, err)
|
|
assert.Len(t, due, 2)
|
|
}
|
|
|
|
type adaptiveGatedHandler struct {
|
|
stubScheduledHandler
|
|
due bool
|
|
err error
|
|
}
|
|
|
|
func (h *adaptiveGatedHandler) ScheduleDue(context.Context, time.Time) (bool, error) {
|
|
return h.due, h.err
|
|
}
|
|
|
|
func TestAdaptiveSchedulerDoesNotCreateEmptyTasks(t *testing.T) {
|
|
truncate(t)
|
|
handler := &adaptiveGatedHandler{stubScheduledHandler: stubScheduledHandler{taskType: "adaptive_gate", enabled: true, interval: 15 * time.Second}}
|
|
withSystemTaskRegistry(t, handler)
|
|
runSystemTaskScheduler()
|
|
assert.Equal(t, int64(0), countSystemTasks(t, handler.taskType))
|
|
handler.due = true
|
|
runSystemTaskScheduler()
|
|
runSystemTaskScheduler()
|
|
assert.Equal(t, int64(1), countSystemTasks(t, handler.taskType))
|
|
}
|