package worker import ( "context" "encoding/json" "errors" "io" "log/slog" "sync" "sync/atomic" "testing" "time" taskbiz "kra/internal/biz/task" ) type workerTaskRepo struct { mu sync.Mutex items []*taskbiz.TimedTask listErr error logs chan *taskbiz.TimedTaskLog } func (r *workerTaskRepo) CreateTask(context.Context, *taskbiz.TimedTask) error { return nil } func (r *workerTaskRepo) UpdateTask(context.Context, *taskbiz.TimedTask) error { return nil } func (r *workerTaskRepo) DeleteTask(context.Context, uint) error { return nil } func (r *workerTaskRepo) ToggleTask(context.Context, uint, bool) error { return nil } func (r *workerTaskRepo) CleanupLogs(context.Context) error { return nil } func (r *workerTaskRepo) TaskNameExists(context.Context, string, uint) (bool, error) { return false, nil } func (r *workerTaskRepo) FindTask(_ context.Context, id uint) (*taskbiz.TimedTask, error) { r.mu.Lock() defer r.mu.Unlock() for _, item := range r.items { if item.ID == id { return cloneTimedTask(item), nil } } return nil, errors.New("task not found") } func (r *workerTaskRepo) ListTasks(context.Context, int, int, *taskbiz.TimedTask) ([]*taskbiz.TimedTask, int64, error) { r.mu.Lock() defer r.mu.Unlock() if r.listErr != nil { return nil, 0, r.listErr } items := make([]*taskbiz.TimedTask, 0, len(r.items)) for _, item := range r.items { items = append(items, cloneTimedTask(item)) } return items, int64(len(items)), nil } func (r *workerTaskRepo) RecordTaskLog(_ context.Context, value *taskbiz.TimedTaskLog) error { if r.logs != nil { copy := *value r.logs <- © } return nil } func (r *workerTaskRepo) ListTaskLogs(context.Context, int, int, uint, string) ([]*taskbiz.TimedTaskLog, int64, error) { return nil, 0, nil } func newTestTaskScheduler(repo *workerTaskRepo) *TaskScheduler { tasks := taskbiz.NewTaskUsecase(repo) executor := &TaskExecutor{tasks: tasks} logger := slog.New(slog.NewTextHandler(io.Discard, nil)) return NewTaskScheduler(tasks, nil, executor, logger) } func TestReloadPreservesExistingScheduleOnPreparationFailure(t *testing.T) { repo := &workerTaskRepo{} scheduler := newTestTaskScheduler(repo) old := &taskbiz.TimedTask{ID: 1, Name: "old", Spec: "0 0 * * *", Enabled: true} if err := scheduler.Schedule(old); err != nil { t.Fatalf("schedule old task: %v", err) } t.Run("query failure", func(t *testing.T) { repo.mu.Lock() repo.listErr = errors.New("database unavailable") repo.mu.Unlock() if err := scheduler.Reload(context.Background()); err == nil { t.Fatal("Reload() error = nil, want query error") } assertOnlyScheduledTask(t, scheduler, old.ID) }) t.Run("invalid cron", func(t *testing.T) { repo.mu.Lock() repo.listErr = nil repo.items = []*taskbiz.TimedTask{{ID: 2, Name: "invalid", Spec: "bad cron", Enabled: true}} repo.mu.Unlock() if err := scheduler.Reload(context.Background()); err == nil { t.Fatal("Reload() error = nil, want cron error") } assertOnlyScheduledTask(t, scheduler, old.ID) }) } func assertOnlyScheduledTask(t *testing.T, scheduler *TaskScheduler, id uint) { t.Helper() scheduler.mu.Lock() defer scheduler.mu.Unlock() if len(scheduler.entries) != 1 { t.Fatalf("scheduled task count = %d, want 1", len(scheduler.entries)) } if _, ok := scheduler.entries[id]; !ok { t.Fatalf("scheduled task %d was replaced after failed reload", id) } if got := len(scheduler.standard.Entries()) + len(scheduler.seconds.Entries()); got != 1 { t.Fatalf("cron entry count = %d, want 1", got) } } func TestConcurrentScheduleKeepsSingleEntry(t *testing.T) { scheduler := newTestTaskScheduler(&workerTaskRepo{}) const workers = 32 var wg sync.WaitGroup for i := 0; i < workers; i++ { wg.Add(1) go func() { defer wg.Done() if err := scheduler.Schedule(&taskbiz.TimedTask{ID: 7, Name: "same", Spec: "0 0 * * *", Enabled: true}); err != nil { t.Errorf("Schedule() error = %v", err) } }() } wg.Wait() assertOnlyScheduledTask(t, scheduler, 7) } func TestManualTriggerSkipsOverlappingExecution(t *testing.T) { const methodName = "worker-test-no-overlap" started := make(chan struct{}, 2) release := make(chan struct{}) var calls atomic.Int32 taskbiz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error { calls.Add(1) started <- struct{}{} <-release return nil }) task := &taskbiz.TimedTask{ID: 9, Name: "single", ExecutorType: taskbiz.TaskExecutorMethod, MethodName: methodName} repo := &workerTaskRepo{items: []*taskbiz.TimedTask{task}, logs: make(chan *taskbiz.TimedTaskLog, 2)} scheduler := newTestTaskScheduler(repo) if err := scheduler.TriggerID(context.Background(), task.ID); err != nil { t.Fatalf("first manual trigger failed: %v", err) } select { case <-started: case <-time.After(time.Second): t.Fatal("first task execution did not start") } if err := scheduler.TriggerID(context.Background(), task.ID); err == nil { t.Fatal("overlapping manual trigger returned success") } if got := calls.Load(); got != 1 { t.Fatalf("execution count while first run is active = %d, want 1", got) } close(release) select { case <-repo.logs: case <-time.After(time.Second): t.Fatal("first task execution did not finish") } waitForSchedulerIdle(t, scheduler) if err := scheduler.TriggerID(context.Background(), task.ID); err != nil { t.Fatalf("manual trigger failed after the previous run finished: %v", err) } select { case <-started: case <-time.After(time.Second): t.Fatal("task could not run again after the first execution finished") } select { case <-repo.logs: case <-time.After(time.Second): t.Fatal("second task execution did not finish") } waitForSchedulerIdle(t, scheduler) if got := calls.Load(); got != 2 { t.Fatalf("total execution count = %d, want 2", got) } } func TestAutomaticAndManualTriggersShareOverlapGate(t *testing.T) { const methodName = "worker-test-auto-manual-overlap" started := make(chan struct{}, 1) release := make(chan struct{}) var calls atomic.Int32 taskbiz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error { calls.Add(1) started <- struct{}{} <-release return nil }) task := &taskbiz.TimedTask{ID: 10, Name: "shared-gate", ExecutorType: taskbiz.TaskExecutorMethod, MethodName: methodName} repo := &workerTaskRepo{logs: make(chan *taskbiz.TimedTaskLog, 1)} scheduler := newTestTaskScheduler(repo) if !scheduler.dispatch(task, "auto", true) { t.Fatal("automatic dispatch was rejected while idle") } select { case <-started: case <-time.After(time.Second): t.Fatal("automatic task execution did not start") } if scheduler.Trigger(task) { t.Fatal("manual trigger bypassed an active automatic execution") } if got := calls.Load(); got != 1 { t.Fatalf("execution count while automatic run is active = %d, want 1", got) } close(release) select { case <-repo.logs: case <-time.After(time.Second): t.Fatal("automatic task execution did not finish") } waitForSchedulerIdle(t, scheduler) } func waitForSchedulerIdle(t *testing.T, scheduler *TaskScheduler) { t.Helper() scheduler.runMu.Lock() idle := scheduler.idle scheduler.runMu.Unlock() select { case <-idle: case <-time.After(time.Second): t.Fatal("scheduler did not become idle") } } func TestStopCancelsAndWaitsForActiveRuns(t *testing.T) { scheduler := newTestTaskScheduler(&workerTaskRepo{}) runCtx, cancel := context.WithCancel(context.Background()) scheduler.ctxMu.Lock() scheduler.runContext = runCtx scheduler.cancel = cancel scheduler.ctxMu.Unlock() taskCtx, ok := scheduler.beginRun(12) if !ok { t.Fatal("beginRun() rejected an idle scheduler") } finished := make(chan struct{}) go func() { <-taskCtx.Done() time.Sleep(20 * time.Millisecond) close(finished) scheduler.finishRun(12) }() stopCtx, stopCancel := context.WithTimeout(context.Background(), time.Second) defer stopCancel() if err := scheduler.Stop(stopCtx); err != nil { t.Fatalf("Stop() error = %v", err) } select { case <-finished: default: t.Fatal("Stop() returned before the active run finished") } if _, ok := scheduler.beginRun(13); ok { t.Fatal("scheduler accepted a new run after Stop()") } }