mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 06:05:45 +08:00
262 lines
8.1 KiB
Go
262 lines
8.1 KiB
Go
package task
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
|
|
"ragflow/internal/entity"
|
|
"ragflow/internal/entity/models"
|
|
)
|
|
|
|
func testStrPtr(s string) *string { return &s }
|
|
|
|
func makeTaskHandlerTestContext(taskType string) *TaskContext {
|
|
var progressCalls []float64
|
|
pipelineID := ""
|
|
if strings.HasPrefix(taskType, "dataflow") {
|
|
pipelineID = "flow-1"
|
|
}
|
|
return &TaskContext{
|
|
IngestionTask: &entity.IngestionTask{
|
|
ID: "task-1",
|
|
DocumentID: "doc-1",
|
|
},
|
|
TaskType: taskType,
|
|
PipelineID: pipelineID,
|
|
Doc: entity.Document{
|
|
ID: "doc-1",
|
|
KbID: "kb-1",
|
|
Name: testStrPtr("test-doc.pdf"),
|
|
ParserID: "naive",
|
|
ParserConfig: entity.JSONMap{},
|
|
},
|
|
KB: entity.Knowledgebase{
|
|
ID: "kb-1",
|
|
TenantID: "tenant-1",
|
|
EmbdID: "embd-1",
|
|
},
|
|
Tenant: entity.Tenant{
|
|
ID: "tenant-1",
|
|
LLMID: "gpt-4",
|
|
},
|
|
ProgressFunc: func(prog float64, msg string) {
|
|
progressCalls = append(progressCalls, prog)
|
|
},
|
|
}
|
|
}
|
|
|
|
func newNoopDataflowService(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) {
|
|
if strings.TrimSpace(dataflowID) == "" {
|
|
dataflowID = "flow-1"
|
|
}
|
|
svc, err := NewDataflowService(ctx, dataflowID, 0, 0)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
svc = svc.
|
|
WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) {
|
|
return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil
|
|
}).
|
|
WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) {
|
|
return map[string]any{
|
|
"chunks": []map[string]any{
|
|
{"text": "stub dataflow chunk", "q_2_vec": []float64{0.1, 0.2}},
|
|
},
|
|
}, dsl, nil
|
|
}).
|
|
WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) {
|
|
return nil, nil
|
|
}).
|
|
WithLogCreateFunc(func(log *entity.PipelineOperationLog) error {
|
|
return nil
|
|
}).
|
|
WithDocService(&stubDocService{}).
|
|
WithChunkCounter(&stubChunkCounter{}).
|
|
WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
|
return nil, nil
|
|
})
|
|
return svc, nil
|
|
}
|
|
|
|
func newNoopTaskHandler(ctx *TaskContext) *TaskHandler {
|
|
return NewTaskHandler(ctx).WithDataflowServiceFactory(newNoopDataflowService)
|
|
}
|
|
|
|
func TestTaskHandler_Dispatch(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
taskType string
|
|
wantErr bool
|
|
wantPanic bool
|
|
}{
|
|
{"memory", "memory", false, false},
|
|
{"dataflow", "dataflow", false, false},
|
|
{"dataflow with suffix", "dataflow_test", false, false},
|
|
{"raptor", "raptor", false, false},
|
|
{"graphrag", "graphrag", false, false},
|
|
{"mindmap", "mindmap", false, false},
|
|
{"evaluation", "evaluation", false, false},
|
|
{"reembedding", "reembedding", false, false},
|
|
{"clone", "clone", false, false},
|
|
{"standard (empty task_type)", "", false, false},
|
|
{"standard (unknown task_type)", "unknown_type", false, false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
ctx := makeTaskHandlerTestContext(tt.taskType)
|
|
handler := newNoopTaskHandler(ctx)
|
|
err := handler.Handle()
|
|
|
|
if tt.wantErr && err == nil {
|
|
t.Error("expected error, got nil")
|
|
}
|
|
if !tt.wantErr && err != nil {
|
|
t.Errorf("unexpected error: %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestTaskHandler_DefaultDataflowServiceInjectsProgress(t *testing.T) {
|
|
ctx := makeTaskHandlerTestContext("dataflow")
|
|
handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) {
|
|
svc, err := NewDataflowService(ctx, dataflowID, 0, 0)
|
|
if err != nil {
|
|
t.Fatalf("NewDataflowService: %v", err)
|
|
}
|
|
if svc.progressFunc == nil {
|
|
t.Fatal("expected default progress func to be injected")
|
|
}
|
|
return newNoopDataflowService(ctx, dataflowID)
|
|
})
|
|
if err := handler.Handle(); err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestTaskHandler_Dataflow_UsesTaskContext(t *testing.T) {
|
|
ctx := makeTaskHandlerTestContext("dataflow")
|
|
type ctxKey string
|
|
const key ctxKey = "trace"
|
|
ctx.Ctx = context.WithValue(context.Background(), key, "task-ctx")
|
|
|
|
handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) {
|
|
return mustNewDataflowService(t, ctx, dataflowID, 0, 0).
|
|
WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) {
|
|
return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil
|
|
}).
|
|
WithRunPipelineFunc(func(runCtx context.Context, dsl string) (map[string]any, string, error) {
|
|
if got := runCtx.Value(key); got != "task-ctx" {
|
|
t.Fatalf("runCtx value = %v, want task-ctx", got)
|
|
}
|
|
return map[string]any{"chunks": []map[string]any{{"text": "stub", "q_2_vec": []float64{0.1, 0.2}}}}, dsl, nil
|
|
}).
|
|
WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) {
|
|
return nil, nil
|
|
}).
|
|
WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }).
|
|
WithDocService(&stubDocService{}).
|
|
WithChunkCounter(&stubChunkCounter{}), nil
|
|
})
|
|
if err := handler.Handle(); err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) {
|
|
ctx := makeTaskHandlerTestContext("dataflow")
|
|
ctx.Doc.PipelineID = testStrPtr("flow-1")
|
|
ctx.Doc.Name = testStrPtr("verify-dataflow.pdf")
|
|
|
|
var pipelineCalled bool
|
|
var insertCalled bool
|
|
var logCreateCalls int
|
|
var insertedChunkCount int
|
|
var progressProgs []float64
|
|
var progressMsgs []string
|
|
|
|
ctx.ProgressFunc = func(prog float64, msg string) {
|
|
progressProgs = append(progressProgs, prog)
|
|
progressMsgs = append(progressMsgs, msg)
|
|
t.Logf("progress: prog=%.2f msg=%q", prog, msg)
|
|
}
|
|
|
|
handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) {
|
|
svc := mustNewDataflowService(t, ctx, dataflowID, 0, 0).
|
|
WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) {
|
|
return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil
|
|
}).
|
|
WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) {
|
|
pipelineCalled = true
|
|
t.Log("mock pipeline.run called")
|
|
return map[string]any{
|
|
"chunks": []map[string]any{
|
|
{"text": "stub dataflow chunk", "q_2_vec": []float64{0.1, 0.2}},
|
|
},
|
|
}, dsl, nil
|
|
}).
|
|
WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) {
|
|
insertCalled = true
|
|
insertedChunkCount = len(chunks)
|
|
t.Logf("mock insertChunks called: %d chunks", len(chunks))
|
|
return nil, nil
|
|
}).
|
|
WithLogCreateFunc(func(log *entity.PipelineOperationLog) error {
|
|
logCreateCalls++
|
|
if log.PipelineID != nil {
|
|
t.Logf("mock pipeline log created: pipeline_id=%s", *log.PipelineID)
|
|
}
|
|
return nil
|
|
}).
|
|
WithDocService(&stubDocService{}).
|
|
WithChunkCounter(&stubChunkCounter{}).
|
|
WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
|
return nil, nil
|
|
})
|
|
return svc, nil
|
|
})
|
|
|
|
if err := handler.Handle(); err != nil {
|
|
t.Fatalf("handler.Handle() error: %v", err)
|
|
}
|
|
|
|
if !pipelineCalled {
|
|
t.Fatal("expected mock pipeline.run to be called")
|
|
}
|
|
if !insertCalled {
|
|
t.Fatal("expected insertChunks to be called")
|
|
}
|
|
if insertedChunkCount != 1 {
|
|
t.Fatalf("insertedChunkCount = %d, want 1", insertedChunkCount)
|
|
}
|
|
if len(progressProgs) == 0 {
|
|
t.Fatal("expected progress callbacks, got none")
|
|
}
|
|
|
|
foundStartIndex := false
|
|
for _, msg := range progressMsgs {
|
|
if strings.Contains(msg, "Start to index") {
|
|
foundStartIndex = true
|
|
break
|
|
}
|
|
}
|
|
if !foundStartIndex {
|
|
t.Fatalf("expected progress message containing %q, got %v", "Start to index", progressMsgs)
|
|
}
|
|
|
|
if got := progressProgs[len(progressProgs)-1]; got != 1.0 {
|
|
t.Fatalf("final progress = %v, want 1.0", got)
|
|
}
|
|
|
|
lastMsg := progressMsgs[len(progressMsgs)-1]
|
|
if !strings.Contains(lastMsg, "Indexing done") {
|
|
t.Fatalf("final progress msg = %q, want substring %q", lastMsg, "Indexing done")
|
|
}
|
|
|
|
if logCreateCalls != 1 {
|
|
t.Fatalf("logCreateCalls = %d, want 1", logCreateCalls)
|
|
}
|
|
}
|