new-api-custom/service/detection_schedule_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/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{Status: "error"}
	if closest != "" {
		result = fingerprint.Result{Status: "inconclusive", Closest: closest, ValidSamples: 1}
	}
	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, "fingerprint", 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, "fingerprint", 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, "fingerprint", 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, "fingerprint", 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, "fingerprint", 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, "fingerprint", time.Unix(6000, 0))
	require.NoError(t, err)
	assert.Equal(t, states, again)
	seedAutomaticFingerprint(t, "abandoned", 7000, "", false, true)
	due, _, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", 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, "fingerprint", 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, "fingerprint", time.Unix(10000, 0))
	require.NoError(t, err)
	assert.Equal(t, states, again)
}

func TestAdaptivePelicanBootstrapInFlightAndRanges(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: "pelican", Status: "succeeded", Quality: "normal", Settled: true, FinishedAt: 1000},
		{ID: "auto2", RequestKey: "b", ChannelID: 2, Kind: "pelican", Status: "failed", Settled: true, FinishedAt: 1000},
		{ID: "paid", RequestKey: "c", UserID: 1, ChannelID: 1, Kind: "pelican", Status: "succeeded", Quality: "degraded", Settled: true, FinishedAt: 1100},
		{ID: "active", RequestKey: "d", ChannelID: 2, Kind: "pelican", Status: "running"},
	}
	require.NoError(t, model.DB.Create(&runs).Error)
	due, states, err := AutomaticDetectionChannelsDue(ctx, "pelican", 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, "pelican")
	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, "fingerprint", 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, "pelican", time.Unix(4799, 0))
	require.NoError(t, err)
	assert.Equal(t, 30, states[2].IntervalMinutes)
	require.Len(t, due, 1)
	due, _, err = AutomaticDetectionChannelsDue(ctx, "pelican", 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))
}