package router import ( "context" "errors" "testing" "time" "github.com/Tencent/WeKnora/internal/types" "github.com/hibiken/asynq" ) func TestSyncTaskExecutorInjectsRetryMetadata(t *testing.T) { executor := NewSyncTaskExecutor() observed := make(chan [2]int, 1) executor.RegisterHandler("test:retry-metadata", func(ctx context.Context, _ *asynq.Task) error { retried, maxRetry, ok := types.TaskRetryMetadataFromContext(ctx) if !ok { observed <- [2]int{-1, -1} return nil } observed <- [2]int{retried, maxRetry} return nil }) task := asynq.NewTask("test:retry-metadata", nil) if _, err := executor.Enqueue(task, asynq.MaxRetry(3)); err != nil { t.Fatalf("enqueue: %v", err) } select { case got := <-observed: if got != [2]int{0, 3} { t.Fatalf("retry metadata = %v, want [0 3]", got) } case <-time.After(2 * time.Second): t.Fatal("timed out waiting for sync task") } } func TestSyncTaskExecutorFencesReservedTaskID(t *testing.T) { executor := NewSyncTaskExecutor() started, release := make(chan struct{}), make(chan struct{}) executor.RegisterHandler("test:reserved-id", func(context.Context, *asynq.Task) error { close(started) <-release return nil }) task := asynq.NewTask("test:reserved-id", []byte(`{"knowledge_id":"document"}`)) opts := []asynq.Option{asynq.TaskID("resumable:document"), asynq.Retention(time.Hour), asynq.MaxRetry(0)} info, err := executor.Enqueue(task, opts...) if err != nil || info.ID != "resumable:document" { t.Fatalf("reserved task: %v %v", info, err) } select { case <-started: case <-time.After(2 * time.Second): t.Fatal("task did not start") } if _, err := executor.Enqueue(task, opts...); !errors.Is(err, asynq.ErrTaskIDConflict) { t.Fatalf("concurrent reserved task was accepted: %v", err) } if !executor.HasActiveTaskID(info.ID) { t.Fatal("active reserved task was not detected") } if !executor.HasActiveKnowledgeTask("document") { t.Fatal("active knowledge task was not detected") } close(release) deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { executor.mu.RLock() expiry := executor.taskIDs[info.ID] executor.mu.RUnlock() if !expiry.IsZero() { if executor.HasActiveTaskID(info.ID) { t.Fatal("retained task reported active") } if executor.HasActiveKnowledgeTask("document") { t.Fatal("completed knowledge task reported active") } if _, err := executor.Enqueue(task, opts...); !errors.Is(err, asynq.ErrTaskIDConflict) { t.Fatalf("retained task was accepted: %v", err) } return } time.Sleep(time.Millisecond) } t.Fatal("completed task was not retained") }