model/system_task_test.go
Rama main · réplica de lectura del repositorio original.
package model
import (
"os"
"testing"
"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 TestSystemTaskHistoryDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var dialector gorm.Dialector
switch dialect {
case "sqlite":
dialector = sqlite.Open(":memory:")
case "mysql":
dsn := os.Getenv("TEST_MYSQL_DSN")
if dsn == "" {
t.Skip("TEST_MYSQL_DSN is not configured")
}
dialector = mysql.Open(dsn)
case "postgres":
dsn := os.Getenv("TEST_POSTGRES_DSN")
if dsn == "" {
t.Skip("TEST_POSTGRES_DSN is not configured")
}
dialector = postgres.Open(dsn)
}
db, err := gorm.Open(dialector, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "history_test_"}})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
previousDB := DB
DB = db
t.Cleanup(func() {
DB = previousDB
require.NoError(t, db.Migrator().DropTable(&SystemTask{}))
require.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&SystemTask{}))
var version string
versionQuery := "SELECT version()"
if dialect == "sqlite" {
versionQuery = "SELECT sqlite_version()"
}
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
t.Logf("database version: %s", version)
tasks := []SystemTask{
{TaskID: "model-failed-old", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusFailed},
{TaskID: "model-success", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusSucceeded},
{TaskID: "channel-failed", Type: SystemTaskTypeChannelTest, Status: SystemTaskStatusFailed},
{TaskID: "model-failed-recent", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusFailed},
{TaskID: "model-latest", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusSucceeded},
{TaskID: "channel-pending", Type: SystemTaskTypeChannelTest, Status: SystemTaskStatusPending},
{TaskID: "channel-running", Type: SystemTaskTypeChannelTest, Status: SystemTaskStatusRunning},
{TaskID: "log-only-history", Type: SystemTaskTypeLogCleanup, Status: SystemTaskStatusSucceeded},
}
require.NoError(t, db.Create(&tasks).Error)
filter := SystemTaskFilter{Scope: "history", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusFailed}
page, total, err := ListSystemTasks(filter, 1, 1)
require.NoError(t, err)
assert.EqualValues(t, 2, total)
require.Len(t, page, 1)
assert.Equal(t, "model-failed-old", page[0].TaskID)
page, total, err = ListSystemTasks(filter, 2, 1)
require.NoError(t, err)
assert.EqualValues(t, 2, total)
assert.Empty(t, page)
active, total, err := ListSystemTasks(SystemTaskFilter{Scope: "active"}, 0, 20)
require.NoError(t, err)
assert.EqualValues(t, 2, total)
require.Len(t, active, 2)
assert.Equal(t, "channel-running", active[0].TaskID)
assert.Equal(t, "channel-pending", active[1].TaskID)
deleted, err := DeleteSystemTaskHistory(filter)
require.NoError(t, err)
assert.EqualValues(t, 2, deleted)
_, total, err = ListSystemTasks(SystemTaskFilter{Scope: "history"}, 0, 20)
require.NoError(t, err)
assert.EqualValues(t, 4, total, "other types and statuses are untouched")
deleted, err = DeleteSystemTaskHistory(SystemTaskFilter{Scope: "active", Status: SystemTaskStatusRunning})
require.NoError(t, err)
assert.Zero(t, deleted, "caller cannot override the terminal-status restriction")
deleted, err = DeleteSystemTaskHistory(SystemTaskFilter{})
require.NoError(t, err)
assert.EqualValues(t, 2, deleted)
remaining, total, err := ListSystemTasks(SystemTaskFilter{}, 0, 100)
require.NoError(t, err)
assert.EqualValues(t, 4, total)
var remainingIDs []string
for _, task := range remaining {
remainingIDs = append(remainingIDs, task.TaskID)
}
assert.ElementsMatch(t, []string{"model-latest", "channel-pending", "channel-running", "log-only-history"}, remainingIDs)
latest, err := GetLatestSystemTasks([]string{SystemTaskTypeModelUpdate, SystemTaskTypeChannelTest, SystemTaskTypeLogCleanup})
require.NoError(t, err)
assert.Equal(t, "model-latest", latest[SystemTaskTypeModelUpdate].TaskID)
assert.Equal(t, "channel-running", latest[SystemTaskTypeChannelTest].TaskID)
assert.Equal(t, "log-only-history", latest[SystemTaskTypeLogCleanup].TaskID)
deleted, err = DeleteSystemTaskHistory(SystemTaskFilter{})
require.NoError(t, err)
assert.Zero(t, deleted, "repeated cleanup preserves the scheduler's latest runs")
})
}
}
type testSystemTaskPayload struct {
TargetTimestamp int64 `json:"target_timestamp"`
BatchSize int `json:"batch_size"`
}
type testSystemTaskState struct {
Total int64 `json:"total"`
Processed int64 `json:"processed"`
Progress int `json:"progress"`
Remaining int64 `json:"remaining"`
}
func createLegacyPendingSystemTask(t *testing.T, taskType string) *SystemTask {
t.Helper()
taskID, err := GenerateSystemTaskID()
require.NoError(t, err)
task := &SystemTask{
TaskID: taskID,
Type: taskType,
Status: SystemTaskStatusPending,
}
require.NoError(t, DB.Create(task).Error)
return task
}
func TestSystemTaskCreateAndActiveLifecycle(t *testing.T) {
truncateTables(t)
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
state := testSystemTaskState{}
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, payload, state)
require.NoError(t, err)
require.NotNil(t, task.ActiveKey)
assert.Equal(t, SystemTaskTypeLogCleanup, *task.ActiveKey)
var decodedPayload testSystemTaskPayload
require.NoError(t, task.DecodePayload(&decodedPayload))
assert.Equal(t, payload, decodedPayload)
activeTask, err := GetActiveSystemTask(SystemTaskTypeLogCleanup)
require.NoError(t, err)
require.NotNil(t, activeTask)
assert.Equal(t, task.TaskID, activeTask.TaskID)
runnerID := "runner-a"
claimedTask, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
err = FinishSystemTask(claimedTask.TaskID, runnerID, SystemTaskStatusSucceeded, map[string]int64{"deleted_count": 0}, "")
require.NoError(t, err)
finishedTask, err := GetSystemTaskByTaskID(task.TaskID)
require.NoError(t, err)
require.NotNil(t, finishedTask)
assert.Nil(t, finishedTask.ActiveKey)
activeTask, err = GetActiveSystemTask(SystemTaskTypeLogCleanup)
require.NoError(t, err)
require.Nil(t, activeTask)
_, err = CreateSystemTask(SystemTaskTypeLogCleanup, payload, state)
require.NoError(t, err)
}
func TestSystemTaskActiveKeyPreventsDuplicateActiveRun(t *testing.T) {
truncateTables(t)
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, payload, testSystemTaskState{})
require.NoError(t, err)
_, err = CreateSystemTask(SystemTaskTypeLogCleanup, payload, testSystemTaskState{})
require.Error(t, err)
activeTask, err := GetActiveSystemTask(SystemTaskTypeLogCleanup)
require.NoError(t, err)
require.NotNil(t, activeTask)
assert.Equal(t, task.TaskID, activeTask.TaskID)
}
func TestSystemTaskLockPreventsConcurrentClaim(t *testing.T) {
truncateTables(t)
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, payload, testSystemTaskState{})
require.NoError(t, err)
secondTask := createLegacyPendingSystemTask(t, SystemTaskTypeLogCleanup)
claimedTask, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-a", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
_, claimed, err = ClaimSystemTask(secondTask.ID, SystemTaskTypeLogCleanup, "runner-b", common.GetTimestamp()+60)
require.NoError(t, err)
require.False(t, claimed)
assert.Equal(t, "runner-a", claimedTask.LockedBy)
reloadedSecond, err := GetSystemTaskByTaskID(secondTask.TaskID)
require.NoError(t, err)
require.NotNil(t, reloadedSecond)
assert.Equal(t, SystemTaskStatusPending, reloadedSecond.Status)
}
func TestExpiredSystemTaskLockFailsOldRunAndClaimsLegacyPendingRun(t *testing.T) {
truncateTables(t)
first, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
_, claimed, err := ClaimSystemTask(first.ID, SystemTaskTypeLogCleanup, "runner-a", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, DB.Model(&SystemTaskLock{}).
Where("task_id = ?", first.TaskID).
Update("locked_until", common.GetTimestamp()-1).Error)
second := createLegacyPendingSystemTask(t, SystemTaskTypeLogCleanup)
claimedTask, claimed, err := ClaimSystemTask(second.ID, SystemTaskTypeLogCleanup, "runner-b", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
assert.Equal(t, second.TaskID, claimedTask.TaskID)
assert.Equal(t, "runner-b", claimedTask.LockedBy)
reloadedFirst, err := GetSystemTaskByTaskID(first.TaskID)
require.NoError(t, err)
require.NotNil(t, reloadedFirst)
assert.Equal(t, SystemTaskStatusFailed, reloadedFirst.Status)
assert.Equal(t, "task lease expired", reloadedFirst.Error)
assert.Nil(t, reloadedFirst.ActiveKey)
}
func TestExpireStaleSystemTaskLockFailsOldRunAndAllowsNewRun(t *testing.T) {
truncateTables(t)
first, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
_, claimed, err := ClaimSystemTask(first.ID, SystemTaskTypeLogCleanup, "runner-a", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, DB.Model(&SystemTaskLock{}).
Where("task_id = ?", first.TaskID).
Update("locked_until", common.GetTimestamp()-1).Error)
require.NoError(t, ExpireStaleSystemTaskLocks(common.GetTimestamp()))
reloadedFirst, err := GetSystemTaskByTaskID(first.TaskID)
require.NoError(t, err)
require.NotNil(t, reloadedFirst)
assert.Equal(t, SystemTaskStatusFailed, reloadedFirst.Status)
assert.Equal(t, "task lease expired", reloadedFirst.Error)
assert.Nil(t, reloadedFirst.ActiveKey)
var lockCount int64
require.NoError(t, DB.Model(&SystemTaskLock{}).Where("task_id = ?", first.TaskID).Count(&lockCount).Error)
assert.Equal(t, int64(0), lockCount)
second, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
require.NotEqual(t, first.TaskID, second.TaskID)
}
func TestFindEarliestPendingSystemTasks(t *testing.T) {
truncateTables(t)
empty, err := FindEarliestPendingSystemTasks(nil)
require.NoError(t, err)
assert.Empty(t, empty)
firstA, err := CreateSystemTask("type_a", nil, nil)
require.NoError(t, err)
ignoredB, err := CreateSystemTask("type_b", nil, nil)
require.NoError(t, err)
_, claimed, err := ClaimSystemTask(ignoredB.ID, "type_b", "runner-b", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, FinishSystemTask(ignoredB.TaskID, "runner-b", SystemTaskStatusFailed, nil, "failed"))
firstB, err := CreateSystemTask("type_b", nil, nil)
require.NoError(t, err)
ignoredC, err := CreateSystemTask("type_c", nil, nil)
require.NoError(t, err)
_, claimed, err = ClaimSystemTask(ignoredC.ID, "type_c", "runner-c", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, FinishSystemTask(ignoredC.TaskID, "runner-c", SystemTaskStatusFailed, nil, "failed"))
tasks, err := FindEarliestPendingSystemTasks([]string{"type_a", "type_b", "type_c", "missing"})
require.NoError(t, err)
require.Len(t, tasks, 2)
assert.Equal(t, firstA.TaskID, tasks["type_a"].TaskID)
assert.Equal(t, firstB.TaskID, tasks["type_b"].TaskID)
assert.Nil(t, tasks["type_c"])
assert.Nil(t, tasks["missing"])
}
func TestGetLatestSystemTask(t *testing.T) {
truncateTables(t)
latest, err := GetLatestSystemTask(SystemTaskTypeChannelTest)
require.NoError(t, err)
require.Nil(t, latest)
first, err := CreateSystemTask(SystemTaskTypeChannelTest, nil, nil)
require.NoError(t, err)
runnerID := "runner-a"
_, claimed, err := ClaimSystemTask(first.ID, SystemTaskTypeChannelTest, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, FinishSystemTask(first.TaskID, runnerID, SystemTaskStatusSucceeded, nil, ""))
second, err := CreateSystemTask(SystemTaskTypeChannelTest, nil, nil)
require.NoError(t, err)
latest, err = GetLatestSystemTask(SystemTaskTypeChannelTest)
require.NoError(t, err)
require.NotNil(t, latest)
assert.Equal(t, second.TaskID, latest.TaskID)
}
func TestGetLatestSystemTasks(t *testing.T) {
truncateTables(t)
empty, err := GetLatestSystemTasks(nil)
require.NoError(t, err)
assert.Empty(t, empty)
firstA, err := CreateSystemTask("type_a", nil, nil)
require.NoError(t, err)
firstB, err := CreateSystemTask("type_b", nil, nil)
require.NoError(t, err)
_, claimed, err := ClaimSystemTask(firstA.ID, "type_a", "runner-a", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, FinishSystemTask(firstA.TaskID, "runner-a", SystemTaskStatusSucceeded, nil, ""))
secondA, err := CreateSystemTask("type_a", nil, nil)
require.NoError(t, err)
tasks, err := GetLatestSystemTasks([]string{"type_a", "type_b", "missing"})
require.NoError(t, err)
require.Len(t, tasks, 2)
assert.NotEqual(t, firstA.TaskID, tasks["type_a"].TaskID)
assert.Equal(t, secondA.TaskID, tasks["type_a"].TaskID)
assert.Equal(t, firstB.TaskID, tasks["type_b"].TaskID)
assert.Nil(t, tasks["missing"])
}
func TestRenewSystemTaskLock(t *testing.T) {
truncateTables(t)
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
runnerID := "runner-a"
_, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
newLockUntil := common.GetTimestamp() + 600
require.NoError(t, RenewSystemTaskLock(task.TaskID, runnerID, newLockUntil))
var lock SystemTaskLock
require.NoError(t, DB.Where("task_id = ?", task.TaskID).First(&lock).Error)
assert.Equal(t, newLockUntil, lock.LockedUntil)
// A different runner cannot renew a lease it does not hold.
assert.ErrorIs(t, RenewSystemTaskLock(task.TaskID, "runner-b", common.GetTimestamp()+600), ErrSystemTaskLockLost)
// After the task finishes it is no longer running, so renew fails.
require.NoError(t, FinishSystemTask(task.TaskID, runnerID, SystemTaskStatusSucceeded, nil, ""))
assert.ErrorIs(t, RenewSystemTaskLock(task.TaskID, runnerID, common.GetTimestamp()+600), ErrSystemTaskLockLost)
}
func TestFinishSystemTaskRetainsExecutor(t *testing.T) {
truncateTables(t)
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
runnerID := "node-1-abc123"
_, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, FinishSystemTask(task.TaskID, runnerID, SystemTaskStatusSucceeded, nil, ""))
reloaded, err := GetSystemTaskByTaskID(task.TaskID)
require.NoError(t, err)
require.NotNil(t, reloaded)
assert.Equal(t, SystemTaskStatusSucceeded, reloaded.Status)
assert.Equal(t, runnerID, reloaded.LockedBy, "executor-of-record must be retained for history")
var lockCount int64
require.NoError(t, DB.Model(&SystemTaskLock{}).Where("task_id = ?", task.TaskID).Count(&lockCount).Error)
assert.Equal(t, int64(0), lockCount)
}
func TestSystemTaskUpdatesRequireCurrentLock(t *testing.T) {
truncateTables(t)
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
runnerID := "runner-a"
_, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, DB.Model(&SystemTaskLock{}).
Where("task_id = ?", task.TaskID).
Updates(map[string]any{"locked_by": "runner-b"}).Error)
assert.ErrorIs(t, UpdateSystemTaskState(task.TaskID, runnerID, testSystemTaskState{Progress: 10}), ErrSystemTaskLockLost)
assert.ErrorIs(t, FinishSystemTask(task.TaskID, runnerID, SystemTaskStatusSucceeded, nil, ""), ErrSystemTaskLockLost)
}
func TestSystemTaskUpdatesRequireUnexpiredLock(t *testing.T) {
truncateTables(t)
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, nil, nil)
require.NoError(t, err)
runnerID := "runner-a"
_, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
require.NoError(t, DB.Model(&SystemTaskLock{}).
Where("task_id = ?", task.TaskID).
Update("locked_until", common.GetTimestamp()-1).Error)
assert.ErrorIs(t, UpdateSystemTaskState(task.TaskID, runnerID, testSystemTaskState{Progress: 10}), ErrSystemTaskLockLost)
assert.ErrorIs(t, FinishSystemTask(task.TaskID, runnerID, SystemTaskStatusSucceeded, nil, ""), ErrSystemTaskLockLost)
reloaded, err := GetSystemTaskByTaskID(task.TaskID)
require.NoError(t, err)
require.NotNil(t, reloaded)
assert.Equal(t, SystemTaskStatusRunning, reloaded.Status)
assert.Empty(t, reloaded.State)
}
func TestUpdateSystemTaskStateIdenticalPayloadDoesNotLoseLock(t *testing.T) {
// SQLite reports matched rows for unchanged UPDATEs, so this case passed
// even before the fix. The MySQL regression is covered by
// TestUpdateSystemTaskStateIdenticalPayloadDoesNotLoseLockConfiguredDatabases.
truncateTables(t)
runUpdateSystemTaskStateIdenticalPayloadKeepsLock(t, SystemTaskTypeLogCleanup)
}
func runUpdateSystemTaskStateIdenticalPayloadKeepsLock(t *testing.T, taskType string) {
t.Helper()
// Two persists in the same second so MySQL's unchanged-row UPDATE returns 0.
task, err := CreateSystemTask(taskType, nil, nil)
require.NoError(t, err)
runnerID := "runner-a"
_, claimed, err := ClaimSystemTask(task.ID, taskType, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
state := testSystemTaskState{Total: 10, Processed: 10, Progress: 100, Remaining: 0}
require.NoError(t, UpdateSystemTaskState(task.TaskID, runnerID, state))
require.NoError(t, UpdateSystemTaskState(task.TaskID, runnerID, state), "identical state persist must not be treated as lock loss")
require.NoError(t, FinishSystemTask(task.TaskID, runnerID, SystemTaskStatusSucceeded, map[string]int64{"deleted_count": 10}, ""))
finished, err := GetSystemTaskByTaskID(task.TaskID)
require.NoError(t, err)
require.NotNil(t, finished)
assert.Equal(t, SystemTaskStatusSucceeded, finished.Status)
}