mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 22:25:41 +08:00
73 lines
2.2 KiB
Go
73 lines
2.2 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/entity"
|
|
taskpkg "ragflow/internal/ingestion/task"
|
|
"ragflow/internal/ingestion/testutil"
|
|
)
|
|
|
|
func TestExecuteTask_DataflowRoutesToTaskHandler(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 routedToDataflow bool
|
|
var gotTaskID string
|
|
var gotProgress []float64
|
|
var gotMsgs []string
|
|
ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error {
|
|
routedToDataflow = true
|
|
gotTaskID = ingestionTask.ID
|
|
wrapped := func(prog float64, msg string) {
|
|
gotProgress = append(gotProgress, prog*100)
|
|
gotMsgs = append(gotMsgs, msg)
|
|
}
|
|
wrapped(0.82, "mock dataflow start")
|
|
wrapped(1.0, "mock dataflow done")
|
|
return nil
|
|
}
|
|
ingestor.storageImpl = testutil.NewMockStorage(map[string][]byte{"/unused": []byte("unused")})
|
|
ingestor.documentDAO = &testutil.MockDocDAO{
|
|
Docs: map[string]*entity.Document{"doc-1": testutil.TestDoc("doc-1", "pdf", ".pdf")},
|
|
}
|
|
|
|
taskCtx := taskpkg.NewTaskContextForScheduling(
|
|
context.Background(),
|
|
&entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING},
|
|
nil,
|
|
)
|
|
|
|
ingestor.executeTask(taskCtx)
|
|
|
|
if !routedToDataflow {
|
|
t.Fatal("expected executeTask to route dataflow task to 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 dataflow done" {
|
|
t.Fatalf("gotMsgs = %v, want final message %q", gotMsgs, "mock dataflow done")
|
|
}
|
|
}
|