WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/router/sync_task_retry_test.go

@@ -0,0 +1,90 @@
+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")
+}