From f409bd756b2cb9def436a3896daed07e3ac099ff Mon Sep 17 00:00:00 2001 From: wizardchen Date: Tue, 14 Jul 2026 15:29:35 +0800 Subject: [PATCH] feat(tests): enhance image multimodal service tests and runtime task handling Added new tests for the ImageMultimodalService to verify the handling of orphan tasks and the finalization of pending counters. Improved runtime task inspection with better error handling and task cancellation logic, ensuring robust management of orphan tasks in the system. --- .../application/service/image_multimodal.go | 3 + .../service/image_multimodal_orphan_test.go | 54 +++++++++++++ internal/handler/system.go | 24 +++--- internal/handler/system_runtime_test.go | 75 ++++++++++++++++++- 4 files changed, 144 insertions(+), 12 deletions(-) diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index 0f00e0121..215d0de53 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -151,6 +151,9 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e logger.Infof(ctx, "[ImageMultimodal] Dropping task chunk=%s knowledge=%s kb=%s image=%s", payload.ChunkID, payload.KnowledgeID, payload.KnowledgeBaseID, payload.ImageURL) + // Still count this image toward the parent finalize gate so a batch + // of dropped orphans cannot strand multimodal:pending forever. + s.checkAndFinalizeAllImages(ctx, payload) return nil } diff --git a/internal/application/service/image_multimodal_orphan_test.go b/internal/application/service/image_multimodal_orphan_test.go index aa7eeb69e..3154dcff1 100644 --- a/internal/application/service/image_multimodal_orphan_test.go +++ b/internal/application/service/image_multimodal_orphan_test.go @@ -9,7 +9,9 @@ import ( "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/alicebob/miniredis/v2" "github.com/hibiken/asynq" + "github.com/redis/go-redis/v9" ) type orphanKnowledgeRepo struct { @@ -96,6 +98,58 @@ func TestImageMultimodalHandleDropsMissingKnowledge(t *testing.T) { } } +type orphanTaskEnqueuer struct { + interfaces.TaskEnqueuer + enqueued []*asynq.Task +} + +func (e *orphanTaskEnqueuer) Enqueue(task *asynq.Task, _ ...asynq.Option) (*asynq.TaskInfo, error) { + e.enqueued = append(e.enqueued, task) + return &asynq.TaskInfo{ID: "post-process"}, nil +} + +func TestImageMultimodalHandleDropFinalizesPendingCounter(t *testing.T) { + t.Parallel() + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("start miniredis: %v", err) + } + t.Cleanup(mr.Close) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + + enqueuer := &orphanTaskEnqueuer{} + svc := &ImageMultimodalService{ + knowledgeRepo: &orphanKnowledgeRepo{err: repository.ErrKnowledgeNotFound}, + kbService: &orphanKBService{kb: &types.KnowledgeBase{ID: "kb-1"}}, + redisClient: rdb, + taskEnqueuer: enqueuer, + } + const knowledgeID = "missing" + redisKey := "multimodal:pending:" + knowledgeID + if err := rdb.Set(context.Background(), redisKey, 1, 0).Err(); err != nil { + t.Fatal(err) + } + payload, err := json.Marshal(types.ImageMultimodalPayload{ + TenantID: 1, + KnowledgeID: knowledgeID, + KnowledgeBaseID: "kb-1", + ImageURL: "minio://bucket/img.png", + }) + if err != nil { + t.Fatal(err) + } + if err := svc.Handle(context.Background(), asynq.NewTask(types.TypeImageMultimodal, payload)); err != nil { + t.Fatalf("orphan drop should succeed: %v", err) + } + if mr.Exists(redisKey) { + t.Fatal("pending counter should be cleared after drop finalize") + } + if len(enqueuer.enqueued) != 1 || enqueuer.enqueued[0].Type() != types.TypeKnowledgePostProcess { + t.Fatalf("expected post-process enqueue, got %d tasks", len(enqueuer.enqueued)) + } +} + func TestImageMultimodalHandlePropagatesTransientKnowledgeError(t *testing.T) { t.Parallel() dbErr := errors.New("db unavailable") diff --git a/internal/handler/system.go b/internal/handler/system.go index 98a31fe62..01d135563 100644 --- a/internal/handler/system.go +++ b/internal/handler/system.go @@ -1827,29 +1827,31 @@ func isKnowledgeGone(err error) bool { } // purgeOrphanRuntimeTask removes queue records for a task whose knowledge row -// is already gone. Best-effort: CancelTasksForKnowledge sweeps every queue -// state; ForceDeleteRuntimeTask is the fallback for the specific task ID. +// is already gone. CancelTasksForKnowledge sweeps sibling tasks best-effort; +// success is reported only once the requested task ID is gone. func (h *SystemHandler) purgeOrphanRuntimeTask( ctx context.Context, inspector interfaces.RuntimeTaskInspector, queue, taskID, knowledgeID string, ) bool { if h.taskInspector != nil && knowledgeID != "" { - deleted, _, err := h.taskInspector.CancelTasksForKnowledge(ctx, knowledgeID) - if err != nil { + if _, _, err := h.taskInspector.CancelTasksForKnowledge(ctx, knowledgeID); err != nil { logger.Warnf(ctx, "purge orphan tasks for knowledge %s: %v", knowledgeID, err) - } else if deleted > 0 { - return true } } supported, err := inspector.ForceDeleteRuntimeTask(ctx, queue, taskID) - if !supported || err != nil { - if err != nil { - logger.Warnf(ctx, "force delete orphan task queue=%s task=%s: %v", queue, taskID, err) - } + if !supported { return false } - return true + if err == nil { + return true + } + // The knowledge-wide sweep may have already removed this task ID. + if _, available, getErr := inspector.GetRuntimeTask(ctx, queue, taskID); available && getErr != nil { + return true + } + logger.Warnf(ctx, "force delete orphan task queue=%s task=%s: %v", queue, taskID, err) + return false } // MutateRuntimeTask executes a backend-advertised action after re-reading the diff --git a/internal/handler/system_runtime_test.go b/internal/handler/system_runtime_test.go index ef1728e9c..664466d79 100644 --- a/internal/handler/system_runtime_test.go +++ b/internal/handler/system_runtime_test.go @@ -3,6 +3,7 @@ package handler import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "testing" @@ -85,6 +86,10 @@ type runtimeTaskTestInspector struct { cancelKnowledge string cancelDeleted int mutatedQueue string + forceDeleteErr error + getRuntimeErr error + getRuntimeErrFrom int + getRuntimeTaskCalls int } func (r *runtimeTaskTestInspector) CancelTasksForKnowledge( @@ -134,6 +139,10 @@ func (r *runtimeTaskTestInspector) ListRuntimeTasks( func (r *runtimeTaskTestInspector) GetRuntimeTask( _ context.Context, queue, taskID string, ) (*types.RuntimeTaskInfo, bool, error) { + r.getRuntimeTaskCalls++ + if r.getRuntimeErr != nil && r.getRuntimeTaskCalls >= r.getRuntimeErrFrom { + return nil, true, r.getRuntimeErr + } for i := range r.tasks { if r.tasks[i].ID == taskID { return &r.tasks[i], true, nil @@ -169,7 +178,7 @@ func (r *runtimeTaskTestInspector) ForceDeleteRuntimeTask( ) (bool, error) { r.mutatedQueue = queue r.forceDeleted = taskID - return true, nil + return true, r.forceDeleteErr } func TestGetRuntimeQueuesReportsIsolatedPoolCapacity(t *testing.T) { @@ -401,3 +410,67 @@ func TestRuntimeTaskCancelPurgesOrphanWhenKnowledgeGone(t *testing.T) { inspector.deletedTask, inspector.forceDeleted, inspector.cancelKnowledge) } } + +func TestRuntimeTaskCancelPurgesOrphanAfterSiblingSweep(t *testing.T) { + gin.SetMode(gin.TestMode) + inspector := &runtimeTaskTestInspector{ + cancelDeleted: 2, + tasks: []types.RuntimeTaskInfo{{ + ID: "task-orphan", Queue: types.QueueMultimodal, Type: types.TypeImageMultimodal, + State: types.RuntimeTaskArchived, TenantID: 42, KnowledgeID: "knowledge-gone", + AllowedActions: []types.RuntimeTaskAction{types.RuntimeTaskActionCancel}, + }}, + } + canceller := &runtimeKnowledgeCancelTest{err: repository.ErrKnowledgeNotFound} + handler := &SystemHandler{taskInspector: inspector, knowledgeSvc: canceller} + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = gin.Params{ + {Key: "queue", Value: types.QueueMultimodal}, + {Key: "task_id", Value: "task-orphan"}, + {Key: "action", Value: string(types.RuntimeTaskActionCancel)}, + } + ctx.Request = httptest.NewRequest(http.MethodPost, "/cancel", nil) + + handler.MutateRuntimeTask(ctx) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + if inspector.cancelKnowledge != "knowledge-gone" || inspector.forceDeleted != "task-orphan" { + t.Fatalf("expected sibling sweep then force delete, got cancel=%q force=%q", + inspector.cancelKnowledge, inspector.forceDeleted) + } +} + +func TestRuntimeTaskCancelPurgesOrphanWhenSweepAlreadyRemovedTask(t *testing.T) { + gin.SetMode(gin.TestMode) + inspector := &runtimeTaskTestInspector{ + cancelDeleted: 1, + forceDeleteErr: errors.New("task not found"), + getRuntimeErr: errors.New("task not found"), + getRuntimeErrFrom: 2, + tasks: []types.RuntimeTaskInfo{{ + ID: "task-orphan", Queue: types.QueueMultimodal, Type: types.TypeImageMultimodal, + State: types.RuntimeTaskRetry, TenantID: 42, KnowledgeID: "knowledge-gone", + AllowedActions: []types.RuntimeTaskAction{types.RuntimeTaskActionCancel}, + }}, + } + canceller := &runtimeKnowledgeCancelTest{err: repository.ErrKnowledgeNotFound} + handler := &SystemHandler{taskInspector: inspector, knowledgeSvc: canceller} + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = gin.Params{ + {Key: "queue", Value: types.QueueMultimodal}, + {Key: "task_id", Value: "task-orphan"}, + {Key: "action", Value: string(types.RuntimeTaskActionCancel)}, + } + ctx.Request = httptest.NewRequest(http.MethodPost, "/cancel", nil) + + handler.MutateRuntimeTask(ctx) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + if inspector.forceDeleted != "task-orphan" { + t.Fatalf("expected force delete attempt, got %q", inspector.forceDeleted) + } +}