feat(store): add TaskStore CRUD and batch query methods for files and blobs
This commit is contained in:
dailz
229
internal/store/task_store_test.go
Normal file
229
internal/store/task_store_test.go
Normal file
@@ -0,0 +1,229 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gcy_hpc_server/internal/model"
|
||||
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
func newTaskTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
|
||||
Logger: logger.Default.LogMode(logger.Silent),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&model.Task{}); err != nil {
|
||||
t.Fatalf("auto migrate: %v", err)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
func makeTestTask(name, status string) *model.Task {
|
||||
return &model.Task{
|
||||
TaskName: name,
|
||||
AppID: 1,
|
||||
AppName: "test-app",
|
||||
Status: status,
|
||||
CurrentStep: "",
|
||||
RetryCount: 0,
|
||||
UserID: "user1",
|
||||
SubmittedAt: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_CreateAndGetByID(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
ctx := context.Background()
|
||||
|
||||
task := makeTestTask("test-task", model.TaskStatusSubmitted)
|
||||
id, err := s.Create(ctx, task)
|
||||
if err != nil {
|
||||
t.Fatalf("Create() error = %v", err)
|
||||
}
|
||||
if id <= 0 {
|
||||
t.Errorf("Create() id = %d, want positive", id)
|
||||
}
|
||||
|
||||
got, err := s.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatalf("GetByID() error = %v", err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatal("GetByID() returned nil")
|
||||
}
|
||||
if got.TaskName != "test-task" {
|
||||
t.Errorf("TaskName = %q, want %q", got.TaskName, "test-task")
|
||||
}
|
||||
if got.Status != model.TaskStatusSubmitted {
|
||||
t.Errorf("Status = %q, want %q", got.Status, model.TaskStatusSubmitted)
|
||||
}
|
||||
if got.ID != id {
|
||||
t.Errorf("ID = %d, want %d", got.ID, id)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_GetByID_NotFound(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
|
||||
got, err := s.GetByID(context.Background(), 999)
|
||||
if err != nil {
|
||||
t.Fatalf("GetByID() error = %v", err)
|
||||
}
|
||||
if got != nil {
|
||||
t.Error("GetByID() expected nil for not-found, got non-nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_ListPagination(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
s.Create(ctx, makeTestTask("task-"+string(rune('A'+i)), model.TaskStatusSubmitted))
|
||||
}
|
||||
|
||||
tasks, total, err := s.List(ctx, &model.TaskListQuery{Page: 1, PageSize: 3})
|
||||
if err != nil {
|
||||
t.Fatalf("List() error = %v", err)
|
||||
}
|
||||
if total != 5 {
|
||||
t.Errorf("total = %d, want 5", total)
|
||||
}
|
||||
if len(tasks) != 3 {
|
||||
t.Errorf("len(tasks) = %d, want 3", len(tasks))
|
||||
}
|
||||
|
||||
tasks2, total2, err := s.List(ctx, &model.TaskListQuery{Page: 2, PageSize: 3})
|
||||
if err != nil {
|
||||
t.Fatalf("List() page 2 error = %v", err)
|
||||
}
|
||||
if total2 != 5 {
|
||||
t.Errorf("total2 = %d, want 5", total2)
|
||||
}
|
||||
if len(tasks2) != 2 {
|
||||
t.Errorf("len(tasks2) = %d, want 2", len(tasks2))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_ListStatusFilter(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
ctx := context.Background()
|
||||
|
||||
s.Create(ctx, makeTestTask("running-1", model.TaskStatusRunning))
|
||||
s.Create(ctx, makeTestTask("running-2", model.TaskStatusRunning))
|
||||
s.Create(ctx, makeTestTask("completed-1", model.TaskStatusCompleted))
|
||||
|
||||
tasks, total, err := s.List(ctx, &model.TaskListQuery{Page: 1, PageSize: 10, Status: model.TaskStatusRunning})
|
||||
if err != nil {
|
||||
t.Fatalf("List() error = %v", err)
|
||||
}
|
||||
if total != 2 {
|
||||
t.Errorf("total = %d, want 2", total)
|
||||
}
|
||||
for _, t2 := range tasks {
|
||||
if t2.Status != model.TaskStatusRunning {
|
||||
t.Errorf("Status = %q, want %q", t2.Status, model.TaskStatusRunning)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_UpdateStatus(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
ctx := context.Background()
|
||||
|
||||
id, _ := s.Create(ctx, makeTestTask("status-test", model.TaskStatusSubmitted))
|
||||
|
||||
err := s.UpdateStatus(ctx, id, model.TaskStatusPreparing, "")
|
||||
if err != nil {
|
||||
t.Fatalf("UpdateStatus() error = %v", err)
|
||||
}
|
||||
|
||||
got, _ := s.GetByID(ctx, id)
|
||||
if got.Status != model.TaskStatusPreparing {
|
||||
t.Errorf("Status = %q, want %q", got.Status, model.TaskStatusPreparing)
|
||||
}
|
||||
if got.StartedAt == nil {
|
||||
t.Error("StartedAt expected non-nil after preparing status")
|
||||
}
|
||||
|
||||
err = s.UpdateStatus(ctx, id, model.TaskStatusFailed, "something broke")
|
||||
if err != nil {
|
||||
t.Fatalf("UpdateStatus(failed) error = %v", err)
|
||||
}
|
||||
got, _ = s.GetByID(ctx, id)
|
||||
if got.ErrorMessage != "something broke" {
|
||||
t.Errorf("ErrorMessage = %q, want %q", got.ErrorMessage, "something broke")
|
||||
}
|
||||
if got.FinishedAt == nil {
|
||||
t.Error("FinishedAt expected non-nil after failed status")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_UpdateRetryState(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
ctx := context.Background()
|
||||
|
||||
task := makeTestTask("retry-test", model.TaskStatusFailed)
|
||||
task.CurrentStep = model.TaskStepDownloading
|
||||
task.RetryCount = 1
|
||||
id, _ := s.Create(ctx, task)
|
||||
|
||||
err := s.UpdateRetryState(ctx, id, model.TaskStatusSubmitted, model.TaskStepDownloading, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("UpdateRetryState() error = %v", err)
|
||||
}
|
||||
|
||||
got, _ := s.GetByID(ctx, id)
|
||||
if got.Status != model.TaskStatusSubmitted {
|
||||
t.Errorf("Status = %q, want %q", got.Status, model.TaskStatusSubmitted)
|
||||
}
|
||||
if got.CurrentStep != model.TaskStepDownloading {
|
||||
t.Errorf("CurrentStep = %q, want %q", got.CurrentStep, model.TaskStepDownloading)
|
||||
}
|
||||
if got.RetryCount != 2 {
|
||||
t.Errorf("RetryCount = %d, want 2", got.RetryCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskStore_GetStuckTasks(t *testing.T) {
|
||||
db := newTaskTestDB(t)
|
||||
s := NewTaskStore(db)
|
||||
ctx := context.Background()
|
||||
|
||||
stuck := makeTestTask("stuck-1", model.TaskStatusDownloading)
|
||||
stuck.UpdatedAt = time.Now().Add(-1 * time.Hour)
|
||||
s.Create(ctx, stuck)
|
||||
|
||||
recent := makeTestTask("recent-1", model.TaskStatusDownloading)
|
||||
s.Create(ctx, recent)
|
||||
|
||||
done := makeTestTask("done-1", model.TaskStatusCompleted)
|
||||
done.UpdatedAt = time.Now().Add(-2 * time.Hour)
|
||||
s.Create(ctx, done)
|
||||
|
||||
tasks, err := s.GetStuckTasks(ctx, 30*time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("GetStuckTasks() error = %v", err)
|
||||
}
|
||||
|
||||
if len(tasks) != 1 {
|
||||
t.Fatalf("len(tasks) = %d, want 1", len(tasks))
|
||||
}
|
||||
if tasks[0].TaskName != "stuck-1" {
|
||||
t.Errorf("TaskName = %q, want %q", tasks[0].TaskName, "stuck-1")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user