internal/ingestion/pipeline/pipeline_test.go
Rama main · réplica de lectura del repositorio original.
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package pipeline
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"testing"
"time"
"ragflow/internal/agent/canvas"
"ragflow/internal/agent/runtime"
"ragflow/internal/entity"
"ragflow/internal/ingestion/component/globals"
"ragflow/internal/parser/media"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"gorm.io/gorm"
)
type mockCanvasStage struct {
output map[string]any
called bool
calls int
}
type pipelineContextMediaCache struct {
scope string
}
func (c *pipelineContextMediaCache) Put(context.Context, string, string, []byte) (entity.MediaRef, error) {
return entity.MediaRef{Scope: c.scope}, nil
}
func (c *pipelineContextMediaCache) Get(context.Context, entity.MediaRef) ([]byte, error) {
return nil, nil
}
type mediaCacheContextStage struct{}
func (*mediaCacheContextStage) Invoke(ctx context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) {
cache, ok := media.CacheFromContext(ctx).(*pipelineContextMediaCache)
if !ok {
return nil, errors.New("pipeline media cache is not attached")
}
if got := globals.TaskID(ctx); got != cache.scope {
return nil, fmt.Errorf("pipeline task id = %q, want media cache scope %q", got, cache.scope)
}
return map[string]any{"media_cache_scope": cache.scope}, nil
}
func (m *mockCanvasStage) Invoke(_ context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) {
m.called = true
m.calls++
out := cloneMapOrEmpty(inputs)
for k, v := range m.output {
out[k] = v
}
return out, nil
}
func (m *mockCanvasStage) Inputs() map[string]string { return map[string]string{"name": "string"} }
func (m *mockCanvasStage) Outputs() map[string]string { return map[string]string{"output": "any"} }
func TestPipelineRunHappyPath(t *testing.T) {
stageA := &mockCanvasStage{output: map[string]any{"a": 1}}
stageB := &mockCanvasStage{output: map[string]any{"b": 2}}
const (
nameA = "p.RunStageA"
nameB = "p.RunStageB"
)
runtime.MustRegister(nameA, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stageA, nil },
runtime.Metadata{Version: "1.0.0"})
runtime.MustRegister(nameB, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stageB, nil },
runtime.Metadata{Version: "1.0.0"})
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
"a": {"obj": {"component_name": "`+nameA+`", "params": {}}, "upstream": ["begin"], "downstream": ["b"]},
"b": {"obj": {"component_name": "`+nameB+`", "params": {}}, "upstream": ["a"]}
},
"path": ["begin", "a", "b"],
"graph": {"nodes": []}
}
}`), "task-canvas-happy")
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
out, err := pipe.Run(t.Context(), map[string]any{"name": "doc-canvas"}, nil)
if err != nil {
t.Fatalf("Run: %v", err)
}
if !stageA.called || !stageB.called {
t.Fatalf("expected both stages to run, got A=%v B=%v", stageA.called, stageB.called)
}
if got := out["name"]; got != "doc-canvas" {
t.Fatalf("name = %v, want doc-canvas", got)
}
gotB, ok := out["b"].(map[string]any)
if !ok {
t.Fatalf("b = %T, want map[string]any", out["b"])
}
if got := gotB["b"]; got != 2 {
t.Fatalf("b.b = %v, want 2", got)
}
}
func TestPipelineRunNilPipeline(t *testing.T) {
var p *Pipeline
if _, err := p.Run(t.Context(), nil, nil); err == nil {
t.Fatal("expected error for nil pipeline")
}
}
func TestPipelineRunStageErrorBubbles(t *testing.T) {
const name = "p.RunErrStage"
runtime.MustRegister(name, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return &errCanvasStage{}, nil },
runtime.Metadata{Version: "1.0.0"})
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["err"]},
"err": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "err"],
"graph": {"nodes": []}
}
}`), "task-canvas-err")
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
if _, err := pipe.Run(t.Context(), map[string]any{"name": "x"}, nil); err == nil {
t.Fatal("expected stage error")
}
}
func TestNewPipelineFromDSLUnwrapsTemplateDSL(t *testing.T) {
pipe, err := NewPipelineFromDSL([]byte(`{
"id": "template-1",
"title": "template",
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}}
},
"path": ["begin"],
"graph": {"nodes": []}
}
}`), "task-template")
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
if pipe.canvas == nil {
t.Fatal("expected decoded canvas")
}
}
type errCanvasStage struct{}
func (e *errCanvasStage) Invoke(_ context.Context, _ *gorm.DB, _ map[string]any) (map[string]any, error) {
return nil, &stageError{Stage: "p.RunErrStage", Reason: "intentional"}
}
func (e *errCanvasStage) Inputs() map[string]string { return nil }
func (e *errCanvasStage) Outputs() map[string]string { return nil }
type factorySentinelStage struct {
marker string
}
func (s *factorySentinelStage) Invoke(_ context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) {
out := cloneMapOrEmpty(inputs)
out["marker"] = s.marker
return out, nil
}
// TestPipelineRun_InstanceFactoryOverridesDefaultFactory verifies that a
// pipeline-scoped component factory can provide task-specific components.
func TestPipelineRun_InstanceFactoryOverridesDefaultFactory(t *testing.T) {
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["stage"]},
"stage": {"obj": {"component_name": "custom-stage", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "stage"],
"graph": {"nodes": []}
}
}`), "task-instance-factory")
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
pipe.WithComponentFactory(func(_ string, _ map[string]any) (runtime.Component, error) {
return &factorySentinelStage{marker: "instance"}, nil
})
out, err := pipe.Run(t.Context(), map[string]any{"name": "doc"}, nil)
if err != nil {
t.Fatalf("Run: %v", err)
}
stage, ok := out["stage"].(map[string]any)
if !ok {
t.Fatalf("stage = %T, want map[string]any", out["stage"])
}
if got := stage["marker"]; got != "instance" {
t.Fatalf("stage.marker = %v, want instance", got)
}
}
func TestPipelineRunAttachesInjectedMediaCacheWithTaskScope(t *testing.T) {
const stageName = "p.MediaCacheContextStage"
runtime.MustRegister(stageName, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) {
return &mediaCacheContextStage{}, nil
},
runtime.Metadata{Version: "1.0.0"})
const taskID = "task-media-cache"
cache := &pipelineContextMediaCache{scope: taskID}
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["stage"]},
"stage": {"obj": {"component_name": "p.MediaCacheContextStage", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "stage"],
"graph": {"nodes": []}
}
}`), taskID, WithDocumentID("doc-media-cache"))
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
pipe.media = cache
out, err := pipe.Run(t.Context(), map[string]any{"name": "image.html"}, nil)
if err != nil {
t.Fatalf("Run: %v", err)
}
stage, ok := out["stage"].(map[string]any)
if !ok {
t.Fatalf("stage output = %T, want map[string]any", out["stage"])
}
if got := stage["media_cache_scope"]; got != taskID {
t.Fatalf("media cache scope = %v, want %q", got, taskID)
}
}
func TestNewPipelineFromDSLRejectsEmptyTaskID(t *testing.T) {
_, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["stage"]},
"stage": {"obj": {"component_name": "Begin", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "stage"],
"graph": {"nodes": []}
}
}`), " ")
if err == nil {
t.Fatal("NewPipelineFromDSL accepted an empty task id")
}
}
func TestPipelineRun_TaskScopedFactoriesDoNotLeakAcrossConcurrentPipelines(t *testing.T) {
origFactory := runtime.DefaultFactory()
runtime.SetDefaultFactory(func(_ string, _ map[string]any) (runtime.Component, error) {
return &factorySentinelStage{marker: "default"}, nil
})
defer runtime.SetDefaultFactory(origFactory)
newPipe := func(taskID string) *Pipeline {
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["stage"]},
"stage": {"obj": {"component_name": "custom-stage", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "stage"],
"graph": {"nodes": []}
}
}`), taskID)
if err != nil {
t.Fatalf("NewPipelineFromDSL(%s): %v", taskID, err)
}
return pipe
}
pipeA := newPipe("task-A")
pipeB := newPipe("task-B")
pipeA.WithComponentFactory(func(_ string, _ map[string]any) (runtime.Component, error) {
return &factorySentinelStage{marker: "A"}, nil
})
pipeB.WithComponentFactory(func(_ string, _ map[string]any) (runtime.Component, error) {
return &factorySentinelStage{marker: "B"}, nil
})
var wg sync.WaitGroup
type result struct {
marker string
err error
}
results := make(chan result, 2)
run := func(pipe *Pipeline) {
defer wg.Done()
out, err := pipe.Run(t.Context(), map[string]any{"name": "doc"}, nil)
if err != nil {
results <- result{err: err}
return
}
stage, ok := out["stage"].(map[string]any)
if !ok {
results <- result{err: fmt.Errorf("stage = %T", out["stage"])}
return
}
results <- result{marker: stage["marker"].(string)}
}
wg.Add(2)
go run(pipeA)
go run(pipeB)
wg.Wait()
close(results)
got := map[string]int{}
for res := range results {
if res.err != nil {
t.Fatalf("Run: %v", res.err)
}
got[res.marker]++
}
if got["A"] != 1 || got["B"] != 1 {
t.Fatalf("markers = %#v, want one A and one B", got)
}
}
// recordingSink captures OnComponentTotal / OnComponentProgress calls so tests
// can assert the pipeline forwards progress to the sink instead of writing
// the DAO layer directly.
type recordingSink struct {
mu sync.Mutex
total int
totalSet bool
events []ProgressEvent
}
func (r *recordingSink) OnComponentTotal(ctx context.Context, taskID string, total int) {
r.mu.Lock()
defer r.mu.Unlock()
r.total = total
r.totalSet = true
}
func (r *recordingSink) OnComponentProgress(ctx context.Context, ev ProgressEvent) {
r.mu.Lock()
defer r.mu.Unlock()
r.events = append(r.events, ev)
}
func TestPipelineTimeoutFailureProgress(t *testing.T) {
const compilerTimeout = "Compiler: Document compilation timed out. Check that the configured model service is reachable and responsive, then retry."
for _, tc := range []struct {
name string
component string
err error
want string
}{
{name: "compiler deadline", component: "Compiler", err: context.DeadlineExceeded, want: compilerTimeout},
{name: "compiler wrapped response deadline", component: "Compiler", err: fmt.Errorf("failed to read response: %w", context.DeadlineExceeded), want: compilerTimeout},
{name: "compiler serialized deadline", component: "Compiler", err: errors.New("failed to read response:context deadline exceeded"), want: compilerTimeout},
{name: "compiler ordinary failure", component: "Compiler", err: errors.New("failed to read response: unexpected EOF"), want: "Compiler: failed to read response: unexpected EOF"},
{name: "compiler canceled", component: "Compiler", err: context.Canceled, want: "Compiler: context canceled"},
{name: "extractor deadline", component: "Extractor", err: context.DeadlineExceeded, want: "Extractor: Document extraction timed out while waiting for the configured LLM. Check that the model service is reachable, then retry."},
} {
t.Run(tc.name, func(t *testing.T) {
sink := &recordingSink{}
const nodeID = "stage:one"
p := &Pipeline{
taskID: "task-timeout", documentID: "doc-timeout", sink: sink,
canvas: &canvas.Canvas{Components: map[string]canvas.CanvasComponent{
nodeID: {Obj: canvas.CanvasComponentObj{ComponentName: tc.component}},
}},
}
p.componentProgressCallback(t.Context())(runtime.ProgressEvent{
Component: nodeID, Phase: runtime.PhaseError, Err: tc.err,
})
if len(sink.events) != 1 {
t.Fatalf("progress events = %d, want 1", len(sink.events))
}
got := sink.events[0]
if got.Message != tc.want || got.Component != nodeID || got.Phase != int(runtime.PhaseError) {
t.Fatalf("failure progress = %+v, want component %q and message %q", got, nodeID, tc.want)
}
})
}
}
// TestPipelineRunForwardsProgressToSink verifies the pipeline reports the
// component-total denominator once via OnComponentTotal and each component
// lifecycle event to the injected ProgressSink.
func TestPipelineRunForwardsProgressToSink(t *testing.T) {
stageA := &mockCanvasStage{output: map[string]any{"a": 1}}
stageB := &mockCanvasStage{output: map[string]any{"b": 2}}
const (
nameA = "p.SinkStageA"
nameB = "p.SinkStageB"
)
runtime.MustRegister(nameA, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stageA, nil },
runtime.Metadata{Version: "1.0.0"})
runtime.MustRegister(nameB, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stageB, nil },
runtime.Metadata{Version: "1.0.0"})
sink := &recordingSink{}
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
"a": {"obj": {"component_name": "`+nameA+`", "params": {}}, "upstream": ["begin"], "downstream": ["b"]},
"b": {"obj": {"component_name": "`+nameB+`", "params": {}}, "upstream": ["a"]}
},
"path": ["begin", "a", "b"],
"graph": {"nodes": []}
}
}`), "task-sink", WithProgressSink(sink), WithDocumentID("doc-sink"))
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
if _, err := pipe.Run(t.Context(), map[string]any{"name": "doc-sink"}, nil); err != nil {
t.Fatalf("Run: %v", err)
}
sink.mu.Lock()
defer sink.mu.Unlock()
if !sink.totalSet || sink.total != 3 {
t.Fatalf("OnComponentTotal = (%d, set=%v), want 3", sink.total, sink.totalSet)
}
if len(sink.events) == 0 {
t.Fatal("expected progress events, got none")
}
seen := map[string]bool{}
for _, ev := range sink.events {
if ev.TaskID != "task-sink" {
t.Fatalf("event TaskID = %q, want task-sink", ev.TaskID)
}
if ev.DocumentID != "doc-sink" {
t.Fatalf("event DocumentID = %q, want doc-sink", ev.DocumentID)
}
seen[ev.Component] = true
}
for _, want := range []string{"a", "b"} {
if !seen[want] {
t.Fatalf("expected progress event for component %q, seen=%v", want, seen)
}
}
}
// =============================================================================
// runPlain — tracker integration with miniredis
// =============================================================================
func TestRunPlain_WithTracker_Success(t *testing.T) {
stage := &mockCanvasStage{output: map[string]any{"result": "ok"}}
const name = "p.RunPlainSuccess"
runtime.MustRegister(name, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stage, nil },
runtime.Metadata{Version: "1.0.0"})
mr := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { client.Close() })
tracker := canvas.NewRunTrackerWithClient(client, time.Hour)
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
"a": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "a"],
"graph": {"nodes": []}
}
}`), "task-tracker-ok", WithRunTracker(tracker))
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
_, err = pipe.Run(t.Context(), map[string]any{"name": "doc"}, nil)
if err != nil {
t.Fatalf("Run: %v", err)
}
}
func TestRunPlain_WithTracker_Error(t *testing.T) {
const name = "p.RunPlainErr"
runtime.MustRegister(name, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return &errCanvasStage{}, nil },
runtime.Metadata{Version: "1.0.0"})
mr := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { client.Close() })
tracker := canvas.NewRunTrackerWithClient(client, time.Hour)
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["err"]},
"err": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "err"],
"graph": {"nodes": []}
}
}`), "task-tracker-err", WithRunTracker(tracker))
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
_, err = pipe.Run(t.Context(), map[string]any{"name": "doc"}, nil)
if err == nil {
t.Fatal("expected stage error, got nil")
}
}
func TestValidatePipeline_DisallowMultipleExtractors(t *testing.T) {
dsl := []byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["Extractor:One"]},
"Extractor:One": {"obj": {"component_name": "Extractor", "params": {}}, "upstream": ["begin"], "downstream": ["Extractor:Two"]},
"Extractor:Two": {"obj": {"component_name": "Extractor", "params": {}}, "upstream": ["Extractor:One"]}
},
"path": ["begin", "Extractor:One", "Extractor:Two"],
"graph": {"nodes": []}
}
}`)
_, err := NewPipelineFromDSL(dsl, "test-task-dup-ext")
if err == nil {
t.Fatal("expected validation error for multiple Extractor components, got nil")
}
if !strings.Contains(err.Error(), "at most 1 Extractor component is allowed") {
t.Errorf("expected at most 1 Extractor error message, got: %v", err)
}
}
// fractionStage reports a fixed in-flight fraction from Invoke, the way a
// parser reports pages or a tokenizer reports embedded chunks.
type fractionStage struct {
fraction float64
}
func (f *fractionStage) Invoke(ctx context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) {
runtime.ReportComponentFraction(ctx, f.fraction)
return cloneMapOrEmpty(inputs), nil
}
func (f *fractionStage) Inputs() map[string]string { return map[string]string{"name": "string"} }
func (f *fractionStage) Outputs() map[string]string { return map[string]string{"output": "any"} }
// fractionRecordingSink records fraction reports alongside the base
// ProgressSink methods, exercising the optional-interface assertion.
type fractionRecordingSink struct {
recordingSink
muFractions sync.Mutex
fractions map[string]float64
}
func (r *fractionRecordingSink) OnComponentFraction(_ context.Context, component string, fraction float64) {
r.muFractions.Lock()
defer r.muFractions.Unlock()
if r.fractions == nil {
r.fractions = map[string]float64{}
}
r.fractions[component] = fraction
}
// TestPipelineRunForwardsFractionsToSink verifies the pipeline assembles the
// fraction channel end to end: a component's bare ReportComponentFraction
// reaches the sink's optional OnComponentFraction under the node's cpnID.
func TestPipelineRunForwardsFractionsToSink(t *testing.T) {
stage := &fractionStage{fraction: 0.42}
const name = "p.FractionStage"
runtime.MustRegister(name, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stage, nil },
runtime.Metadata{Version: "1.0.0"})
sink := &fractionRecordingSink{}
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
"a": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "a"],
"graph": {"nodes": []}
}
}`), "task-frac", WithProgressSink(sink), WithDocumentID("doc-frac"))
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
if _, err := pipe.Run(t.Context(), map[string]any{"name": "doc-frac"}, nil); err != nil {
t.Fatalf("Run: %v", err)
}
sink.muFractions.Lock()
defer sink.muFractions.Unlock()
if got := sink.fractions["a"]; got != 0.42 {
t.Fatalf("fraction for node a = %v, want 0.42 (fractions=%v)", got, sink.fractions)
}
}
// TestPipelineRunWithoutFractionSink verifies a sink that does not implement
// the optional fraction interface still runs: the channel is simply absent.
func TestPipelineRunWithoutFractionSink(t *testing.T) {
stage := &fractionStage{fraction: 0.9}
const name = "p.FractionStageNoSink"
runtime.MustRegister(name, runtime.CategoryIngestion,
func(_ string, _ map[string]any) (runtime.Component, error) { return stage, nil },
runtime.Metadata{Version: "1.0.0"})
pipe, err := NewPipelineFromDSL([]byte(`{
"dsl": {
"components": {
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
"a": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
},
"path": ["begin", "a"],
"graph": {"nodes": []}
}
}`), "task-nofrac", WithProgressSink(&recordingSink{}))
if err != nil {
t.Fatalf("NewPipelineFromDSL: %v", err)
}
if _, err := pipe.Run(t.Context(), map[string]any{"name": "x"}, nil); err != nil {
t.Fatalf("Run: %v", err)
}
}