new-api-custom/service/astra_sampling_test.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{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{Status: status, ValidSamples: 3, Closest: closest}})
	}
	rounds = append(rounds, AstraSamplingRound{ChannelID: 99, StartedAt: now.Unix(), Result: fingerprint.Result{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 TestAstraSamplingAggregationCountsLunaLikeInconclusiveWithinSixHours(t *testing.T) {
	now := time.Now()
	channel := model.Channel{Id: 1, Group: "example-variable", Models: fingerprint.Model}
	rounds := []AstraSamplingRound{
		{ID: "recent-match", ChannelID: 1, StartedAt: now.Add(-time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{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{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{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)["example-variable"]
	if summary.Rate == nil || *summary.Rate != .5 || summary.Matches != 1 || summary.Mismatches != 1 || summary.Inconclusive != 0 {
		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) {
	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 TestAstraSamplingCountsAllValidNonAstraFingerprints(t *testing.T) {
	now := time.Now()
	results := []fingerprint.Result{
		{Status: "error", ValidSamples: 1, Closest: "gpt-5.6-luna", Gap: .99},
		{Status: "inconclusive", ValidSamples: 3, Closest: "gpt-5.5", Gap: .17},
		{Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-sol", Gap: .09},
		{Status: "error", ValidSamples: 0},
	}
	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, 1.0, *summary.Rate)
	assert.Equal(t, 3, summary.Mismatches)
	assert.Equal(t, 0, summary.Inconclusive)
	assert.Equal(t, 1, summary.Errors)
	for _, h := range summary.History {
		if h.ValidSamples > 0 {
			assert.Equal(t, "mismatch", h.Status)
		}
	}
	assert.Equal(t, "match", AstraSamplingResultStatus(fingerprint.Result{Status: "inconclusive", ValidSamples: 1, Closest: fingerprint.Model}))
	assert.Equal(t, "inconclusive", AstraSamplingResultStatus(fingerprint.Result{Status: "mismatch", ValidSamples: 0, Closest: "gpt-5.6-sol"}))
	assert.Equal(t, "running", AstraSamplingResultStatus(fingerprint.Result{Status: "running", ValidSamples: 1, Closest: "gpt-5.6-sol"}))
}