84 lines
2.7 KiB
Go
84 lines
2.7 KiB
Go
package model
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
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 TestSystemTaskActiveKeyIsReleasedOnFinish(t *testing.T) {
|
|
truncateTables(t)
|
|
|
|
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
|
|
state := testSystemTaskState{}
|
|
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, state)
|
|
require.NoError(t, err)
|
|
|
|
var decodedPayload testSystemTaskPayload
|
|
require.NoError(t, task.DecodePayload(&decodedPayload))
|
|
assert.Equal(t, payload, decodedPayload)
|
|
|
|
_, err = CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, state)
|
|
require.Error(t, err)
|
|
|
|
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)
|
|
|
|
activeTask, err = GetActiveSystemTask(SystemTaskTypeLogCleanup)
|
|
require.NoError(t, err)
|
|
require.Nil(t, activeTask)
|
|
|
|
_, err = CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, state)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func TestSystemTaskClaimRequiresExpiredLock(t *testing.T) {
|
|
truncateTables(t)
|
|
|
|
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
|
|
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, testSystemTaskState{})
|
|
require.NoError(t, err)
|
|
|
|
claimedTask, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-a", common.GetTimestamp()+60)
|
|
require.NoError(t, err)
|
|
require.True(t, claimed)
|
|
|
|
_, claimed, err = ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-b", common.GetTimestamp()+60)
|
|
require.NoError(t, err)
|
|
require.False(t, claimed)
|
|
|
|
require.NoError(t, DB.Model(claimedTask).Updates(map[string]any{
|
|
"locked_until": common.GetTimestamp() - 1,
|
|
}).Error)
|
|
|
|
claimedTask, claimed, err = ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-b", common.GetTimestamp()+60)
|
|
require.NoError(t, err)
|
|
require.True(t, claimed)
|
|
assert.Equal(t, "runner-b", claimedTask.LockedBy)
|
|
}
|