155 lines
7.6 KiB
Go
155 lines
7.6 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/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type anchoredSamplingTestHandler struct{ stubScheduledHandler }
|
|
|
|
func (*anchoredSamplingTestHandler) ScheduleFromCreation() bool { return true }
|
|
|
|
func TestAstraSamplingThirtyMinuteSchedulePersistsAcrossRestart(t *testing.T) {
|
|
truncate(t)
|
|
handler := &anchoredSamplingTestHandler{stubScheduledHandler{taskType: AstraSamplingTaskType, enabled: true, interval: 30 * time.Minute}}
|
|
withSystemTaskRegistry(t, handler)
|
|
now := common.GetTimestamp()
|
|
state, _ := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ChannelID: 42, StartedAt: now - 1900, FinishedAt: now - 1800, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}}}})
|
|
row := model.SystemTask{TaskID: "sampling-old", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 1900, UpdatedAt: now - 10, State: string(state)}
|
|
if err := model.DB.Create(&row).Error; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
runSystemTaskScheduler()
|
|
if countSystemTasks(t, AstraSamplingTaskType) != 2 {
|
|
t.Fatal("schedule drifted from completion instead of creation")
|
|
}
|
|
runSystemTaskScheduler()
|
|
if countSystemTasks(t, AstraSamplingTaskType) != 2 {
|
|
t.Fatal("restart duplicated pending job")
|
|
}
|
|
channels, err := AstraSamplingChannels(context.Background())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_ = channels
|
|
loaded, err := model.GetSystemTaskByTaskID(row.TaskID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
summary := AggregateAstraSampling([]model.Channel{{Id: 42, Group: "group", Models: fingerprint.Model}}, []*model.SystemTask{loaded}, time.Now())
|
|
if summary["group"].Matches != 1 {
|
|
t.Fatal("completed results did not survive reload")
|
|
}
|
|
}
|
|
|
|
func TestAstraSamplingAggregationPrivacyAndFailures(t *testing.T) {
|
|
now := time.Now()
|
|
channels := []model.Channel{{Id: 1, Group: "first,shared", Models: fingerprint.Model}, {Id: 2, Group: "first", Models: fingerprint.Model}, {Id: 3, Group: "empty", Models: fingerprint.Model}}
|
|
rounds := []AstraSamplingRound{}
|
|
for _, status := range []string{"match", "mismatch", "inconclusive", "error", "running"} {
|
|
closest := ""
|
|
if status == "match" {
|
|
closest = fingerprint.Model
|
|
}
|
|
if status == "mismatch" {
|
|
closest = "gpt-5.6-luna"
|
|
}
|
|
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: status, ValidSamples: 3, Closest: closest}})
|
|
}
|
|
rounds = append(rounds, AstraSamplingRound{ChannelID: 99, StartedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch"}})
|
|
state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
summaries := AggregateAstraSampling(channels, []*model.SystemTask{{Status: model.SystemTaskStatusFailed, State: string(state)}}, now)
|
|
first := summaries["first"]
|
|
if first.Rate == nil || *first.Rate != .5 || first.Errors != 2 || first.Inconclusive != 1 || first.Running != 0 || first.SampledChannels != 1 || first.Channels != 2 {
|
|
t.Fatalf("%+v", first)
|
|
}
|
|
if *summaries["shared"].Rate != .5 || summaries["empty"].Rate != nil {
|
|
t.Fatal("shared results / empty denominator")
|
|
}
|
|
public, _ := common.Marshal(summaries)
|
|
var raw map[string]any
|
|
if common.Unmarshal(public, &raw) != nil {
|
|
t.Fatal("invalid public JSON")
|
|
}
|
|
for _, h := range first.History {
|
|
if h.Status == "running" {
|
|
t.Fatal("restart must fail unfinished sample")
|
|
}
|
|
}
|
|
if len(first.History) != 5 {
|
|
t.Fatal("deleted/unrelated channel leaked into current group")
|
|
}
|
|
}
|
|
|
|
func TestAstraSamplingExcludesInconclusiveFromRateWithinSixHours(t *testing.T) {
|
|
now := time.Now()
|
|
channel := model.Channel{Id: 1, Group: "gpt-s-unstable", Models: fingerprint.Model}
|
|
rounds := []AstraSamplingRound{
|
|
{ID: "recent-match", ChannelID: 1, StartedAt: now.Add(-time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}},
|
|
{ID: "recent-luna", ChannelID: 1, StartedAt: now.Add(-2 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
|
|
{ID: "old-luna", ChannelID: 1, StartedAt: now.Add(-7 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
|
|
}
|
|
state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
summary := AggregateAstraSampling([]model.Channel{channel}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["gpt-s-unstable"]
|
|
if summary.Rate == nil || *summary.Rate != 0 || summary.Matches != 1 || summary.Mismatches != 0 || summary.Inconclusive != 1 {
|
|
t.Fatalf("%+v", summary)
|
|
}
|
|
if summary.WindowHours != 6 || summary.IntervalMinutes != 30 || len(summary.History) != 2 {
|
|
t.Fatalf("unexpected window, interval or history: %+v", summary)
|
|
}
|
|
}
|
|
|
|
func TestAstraSamplingAggregationMarksNonAstraGroupsAvailabilityOnly(t *testing.T) {
|
|
t.Setenv("CHANNEL_AVAILABILITY_INTERVAL_MINUTES", "30")
|
|
now := time.Now()
|
|
channels := []model.Channel{{Id: 1, Group: "opencode", Models: "gpt-5.6-sol"}}
|
|
summary := AggregateAstraSampling(channels, nil, now)["opencode"]
|
|
if !summary.AvailabilityOnly || summary.Channels != 0 || summary.IntervalMinutes != 30 || summary.Rate != nil {
|
|
t.Fatalf("%+v", summary)
|
|
}
|
|
state, _ := common.Marshal(ChannelAvailabilityState{Rounds: []ChannelAvailabilityRound{
|
|
{ChannelID: 1, At: now.Add(-time.Hour).Unix(), Success: true},
|
|
{ChannelID: 1, At: now.Add(-2 * time.Hour).Unix(), Success: false},
|
|
{ChannelID: 1, At: now.Add(-7 * time.Hour).Unix(), Success: false},
|
|
}})
|
|
AggregateChannelAvailability(map[string]*AstraSamplingSummary{"opencode": summary}, channels, []*model.SystemTask{{State: string(state)}}, now)
|
|
if summary.AvailabilityRate == nil || *summary.AvailabilityRate != .5 || summary.AvailabilityChecks != 2 || summary.AvailabilitySuccesses != 1 {
|
|
t.Fatalf("%+v", summary)
|
|
}
|
|
}
|
|
|
|
func TestAstraSamplingExcludesOldEngineAndUncertainCandidates(t *testing.T) {
|
|
now := time.Now()
|
|
results := []fingerprint.Result{
|
|
{Engine: "meow-4.5.4", Status: "mismatch", ValidSamples: 32, Closest: "other"},
|
|
{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "other"},
|
|
{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model},
|
|
}
|
|
var rounds []AstraSamplingRound
|
|
for _, result := range results {
|
|
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: result})
|
|
}
|
|
state, err := common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
require.NoError(t, err)
|
|
summary := AggregateAstraSampling([]model.Channel{{Id: 1, Group: "test", Models: fingerprint.Model}}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["test"]
|
|
require.NotNil(t, summary.Rate)
|
|
assert.Equal(t, 0.0, *summary.Rate)
|
|
assert.Equal(t, 1, summary.Matches)
|
|
assert.Equal(t, 0, summary.Mismatches)
|
|
assert.Equal(t, 1, summary.Inconclusive)
|
|
assert.Len(t, summary.History, 2)
|
|
}
|
|
|
|
func TestModelTraceVerdictDoesNotPromoteClosestCandidate(t *testing.T) {
|
|
for _, status := range []string{"match", "mismatch", "inconclusive", "error"} {
|
|
assert.Equal(t, status, AstraSamplingResultStatus(fingerprint.Result{Engine: fingerprint.Version, Status: status, ValidSamples: 32, Closest: fingerprint.Model, Weight: .63436}))
|
|
}
|
|
}
|