ragflow/internal/ingestion/service/execute_task_test.go

package service

import (
	"context"
	"testing"

	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/entity"
	taskpkg "ragflow/internal/ingestion/task"
	"ragflow/internal/ingestion/testutil"
)

// TestExecuteTask_CheckpointParseFailureDoesNotKillProcess verifies that checkpoint
// parse failures do not call fatal exit (which would kill the whole worker process).
// Instead, the task should be marked as FAILED and return gracefully.
// This tests the fix for issue 1 from the code review.
func TestExecuteTask_CheckpointParseFailureDoesNotKillProcess(t *testing.T) {
	db := testutil.SetupTestDB(t)
	cleanup := testutil.ReplaceDBForTest(t, db)
	defer cleanup()

	_, _, docID, taskID := testutil.SeedTestData(t, db,
		testutil.WithPipelineID("flow-1"),
		testutil.WithTenantID("tenant-1"),
	)

	// Create a task log with invalid checkpoint (run_count is a string instead of number)
	err := db.Create(&entity.IngestionTaskLog{
		TaskID: taskID,
		Checkpoint: entity.JSONMap{
			"run_count": "not-a-number", // intentionally wrong type
		},
	}).Error
	if err != nil {
		t.Fatalf("create bad task log: %v", err)
	}

	ingestor := NewIngestor("test", 1, []string{"pdf"})
	// Replace runDocumentTask to ensure it doesn't get called
	var runDocumentTaskCalled bool
	ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error {
		runDocumentTaskCalled = true
		return nil
	}

	taskCtx := taskpkg.NewTaskContextForScheduling(
		context.Background(),
		&entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING},
	)

	// Execute the task - this should NOT panic or fatal exit (this is our main validation!)
	ingestor.executeTask(taskCtx)

	// Corrupted run_count values are skipped by IncrementRunCount, so the task
	// proceeds to runDocumentTask and completes normally.
	if !runDocumentTaskCalled {
		t.Fatal("expected runDocumentTask to be called (bad run_count is skipped, not fatal)")
	}

	// Verify task status was set to COMPLETED
	finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID)
	if err != nil {
		t.Fatalf("load final ingestion task: %v", err)
	}
	if finalTask.Status != common.COMPLETED {
		t.Fatalf("final status = %s, want %s", finalTask.Status, common.COMPLETED)
	}
}

func TestDefaultRunDocumentTask_RequiresConfiguredPipelineID(t *testing.T) {
	db := testutil.SetupTestDB(t)
	cleanup := testutil.ReplaceDBForTest(t, db)
	defer cleanup()

	_, kbID, docID, taskID := testutil.SeedTestData(t, db,
		testutil.WithTenantID("tenant-1"),
		testutil.WithKBID("kb-1"),
		testutil.WithDocID("doc-1"),
		testutil.WithTaskID("task-1"),
	)

	ingestor := NewIngestor("test", 1, []string{"pdf"})
	err := ingestor.defaultRunDocumentTask(context.Background(), &entity.IngestionTask{
		ID:         taskID,
		DocumentID: docID,
		DatasetID:  kbID,
		Status:     common.RUNNING,
	})
	if err == nil {
		t.Fatal("expected error when no pipeline is configured")
	}
	if err.Error() != "ingestion task task-1: no pipeline_id configured for document doc-1 or dataset kb-1" {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestExecuteTask_RunsDocumentTask(t *testing.T) {
	db := testutil.SetupTestDB(t)
	cleanup := testutil.ReplaceDBForTest(t, db)
	defer cleanup()

	_, _, docID, taskID := testutil.SeedTestData(t, db,
		testutil.WithPipelineID("flow-1"),
		testutil.WithTenantID("tenant-1"),
	)

	ingestor := NewIngestor("test", 1, []string{"pdf"})
	var runDocumentTaskCalled bool
	var gotTaskID string
	var gotProgress []float64
	var gotMsgs []string
	ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error {
		runDocumentTaskCalled = true
		gotTaskID = ingestionTask.ID
		wrapped := func(prog float64, msg string) {
			gotProgress = append(gotProgress, prog*100)
			gotMsgs = append(gotMsgs, msg)
		}
		wrapped(0.82, "mock pipeline start")
		wrapped(1.0, "mock pipeline done")
		return nil
	}

	taskCtx := taskpkg.NewTaskContextForScheduling(
		context.Background(),
		&entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING},
	)

	ingestor.executeTask(taskCtx)

	if !runDocumentTaskCalled {
		t.Fatal("expected executeTask to run runDocumentTask")
	}
	if gotTaskID != taskID {
		t.Fatalf("runDocumentTask got task ID %q, want %q", gotTaskID, taskID)
	}
	finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID)
	if err != nil {
		t.Fatalf("load final ingestion task: %v", err)
	}
	if finalTask.Status != common.COMPLETED {
		t.Fatalf("final status = %s, want %s", finalTask.Status, common.COMPLETED)
	}
	if len(gotProgress) != 2 || gotProgress[0] != 82 || gotProgress[1] != 100 {
		t.Fatalf("gotProgress = %v, want [82 100]", gotProgress)
	}
	if len(gotMsgs) != 2 || gotMsgs[1] != "mock pipeline done" {
		t.Fatalf("gotMsgs = %v, want final message %q", gotMsgs, "mock pipeline done")
	}
}

// TestExecuteTask_CancelBeforePipeline verifies that when cancelCheck returns
// true at task start, the task is cancelled before AdvanceStep,
// runDocumentTask is never called, and document progress is set to -1 with a
// cancel marker. Mirrors Python's cancel flow where has_canceled() returns
// true in Pipeline.callback().
func TestExecuteTask_CancelBeforePipeline(t *testing.T) {
	db := testutil.SetupTestDB(t)
	cleanup := testutil.ReplaceDBForTest(t, db)
	defer cleanup()
	_, _, docID, taskID := testutil.SeedTestData(t, db,
		testutil.WithPipelineID("flow-1"),
		testutil.WithTenantID("tenant-1"),
	)

	ingestor := NewIngestor("test", 1, []string{"pdf"})
	ingestor.cancelCheck = func(taskID string) bool { return true }

	var runDocumentTaskCalled bool
	ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error {
		runDocumentTaskCalled = true
		return nil
	}

	taskCtx := taskpkg.NewTaskContextForScheduling(
		context.Background(),
		&entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING},
	)
	ingestor.executeTask(taskCtx)

	if runDocumentTaskCalled {
		t.Fatal("expected runDocumentTask to NOT be called when cancel is detected before pipeline")
	}

	doc, err := dao.NewDocumentDAO().GetByID(docID)
	if err != nil {
		t.Fatalf("load document: %v", err)
	}
	if doc.Progress != -1 {
		t.Fatalf("document.progress = %v, want -1 (cancelled)", doc.Progress)
	}
	if doc.Run == nil || *doc.Run != string(entity.TaskStatusCancel) {
		t.Fatalf("document.run = %v, want %s (CANCEL)", doc.Run, entity.TaskStatusCancel)
	}
	if doc.ProgressMsg == nil || *doc.ProgressMsg == "" {
		t.Fatal("document.progress_msg should contain cancel marker, got empty")
	}
}