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) }) } }