new-api-custom/model/detection_schedule_test.go

package model

import (
	"context"
	"fmt"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"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"
)

func TestAdaptiveDetectionIntervals(t *testing.T) {
	state := DetectionSchedule{IntervalMinutes: 30}
	for _, step := range []struct {
		result   string
		interval int
	}{
		{"normal", 30}, {"normal", 60}, {"normal", 60}, {"degraded", 15}, {"normal", 15}, {"normal", 30}, {"normal", 60},
		{"unknown", 15}, {"unknown", 15}, {"normal", 30}, {"normal", 60}, {"failed", 15}, {"degraded", 15}, {"degraded", 30}, {"degraded", 60},
	} {
		state.IntervalMinutes = NextDetectionInterval(state.LastValid, step.result, state.IntervalMinutes)
		assert.Equal(t, step.interval, state.IntervalMinutes, step.result)
		if step.result == "normal" || step.result == "degraded" {
			state.LastValid = step.result
		}
	}
	assert.Equal(t, 15, NextDetectionInterval("", "unknown", 30))
	assert.Equal(t, 30, NextDetectionInterval("", "degraded", 15))
}

func TestAdaptiveDetectionDatabaseMatrix(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("adaptive_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			conn, err := db.DB()
			require.NoError(t, err)
			conn.SetMaxOpenConns(1)
			if dialect != "sqlite" {
				conn.SetMaxOpenConns(4)
			}
			old := DB
			DB = db
			oldType := common.MainDatabaseType()
			common.SetMainDatabaseType(common.DatabaseType(dialect))
			t.Cleanup(func() { common.SetMainDatabaseType(oldType) })
			t.Cleanup(func() {
				for _, row := range []any{&DetectionSchedule{}, &DetectionRun{}, &SystemTask{}} {
					require.NoError(t, db.Migrator().DropTable(row))
				}
				DB = old
				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.Logf("database version: %s", version)
			// Fresh schema, then the exact pre-upgrade task models with real rows.
			require.NoError(t, db.AutoMigrate(&DetectionSchedule{}))
			require.NoError(t, db.Migrator().DropTable(&DetectionSchedule{}))
			require.NoError(t, db.AutoMigrate(&DetectionRun{}, &SystemTask{}))
			oldRun := DetectionRun{ID: "old", RequestKey: "old-unique", ChannelID: 12, Kind: "pelican", Status: "succeeded", Quality: "normal", Settled: true, FinishedAt: 100, Charged: 17}
			require.NoError(t, db.Create(&oldRun).Error)
			require.NoError(t, db.Create(&SystemTask{TaskID: "old-task", Type: "astra_fingerprint", State: `{"rounds":[]}`}).Error)
			for range 2 {
				require.NoError(t, db.AutoMigrate(&DetectionRun{}, &SystemTask{}, &DetectionSchedule{}))
			}
			var retained DetectionRun
			require.NoError(t, db.First(&retained, "id = ?", "old").Error)
			assert.Equal(t, 17, retained.Charged)
			duplicate := oldRun
			duplicate.ID = "duplicate"
			require.Error(t, db.Create(&duplicate).Error)
			ctx := context.Background()
			seed := DetectionSchedule{ChannelID: 12, Kind: "pelican", IntervalMinutes: 30}
			require.NoError(t, EnsureDetectionSchedule(ctx, seed))
			require.NoError(t, EnsureDetectionSchedule(ctx, seed))
			require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "pelican", "one", 100, "normal"))
			require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "pelican", "two", 1900, "normal"))
			var wg sync.WaitGroup
			errs := make(chan error, 2)
			for range 2 {
				wg.Go(func() { errs <- ApplyDetectionScheduleResult(ctx, 12, "pelican", "change", 5500, "degraded") })
			}
			wg.Wait()
			close(errs)
			for err := range errs {
				require.NoError(t, err)
			}
			require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "pelican", "stale", 1901, "normal"))
			var loaded DetectionSchedule
			require.NoError(t, db.First(&loaded, "channel_id = ? AND kind = ?", 12, "pelican").Error)
			assert.Equal(t, 15, loaded.IntervalMinutes)
			assert.Equal(t, int64(6400), loaded.NextDueAt)
			assert.Equal(t, "degraded", loaded.LastValid)
			seed.Kind = "fingerprint"
			require.NoError(t, EnsureDetectionSchedule(ctx, seed))
			require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "fingerprint", "fp", 5600, "normal"))
			require.NoError(t, db.Migrator().DropTable(&SystemTask{}))
			require.NoError(t, db.AutoMigrate(&SystemTask{}, &DetectionSchedule{}))
			var rows []DetectionSchedule
			require.NoError(t, db.Find(&rows).Error)
			require.Len(t, rows, 2)
			// Reinitializing from history must not overwrite a persisted clock.
			seed.Kind = "pelican"
			require.NoError(t, EnsureDetectionSchedule(ctx, seed))
			require.NoError(t, db.First(&loaded, "channel_id = ? AND kind = ?", 12, "pelican").Error)
			assert.Equal(t, int64(6400), loaded.NextDueAt)
		})
	}
}