mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-10 21:55:42 +08:00
371 lines
11 KiB
Go
371 lines
11 KiB
Go
//
|
|
// 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 task
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/engine"
|
|
"ragflow/internal/engine/elasticsearch"
|
|
"ragflow/internal/engine/infinity"
|
|
"ragflow/internal/entity"
|
|
"ragflow/internal/ingestion/testutil"
|
|
"ragflow/internal/server"
|
|
)
|
|
|
|
// =============================================================================
|
|
// Test helper: setupTestDocEngine - Creates either Elasticsearch or Infinity engine
|
|
// =============================================================================
|
|
|
|
func setupTestDocEngine(t *testing.T, engineType engine.EngineType, tenantID, datasetID string) (engine.DocEngine, func()) {
|
|
t.Helper()
|
|
|
|
var (
|
|
docEngine engine.DocEngine
|
|
err error
|
|
)
|
|
|
|
switch engineType {
|
|
case engine.EngineElasticsearch:
|
|
t.Logf("Setting up Elasticsearch engine...")
|
|
esHost := os.Getenv("ES_HOST")
|
|
if esHost == "" {
|
|
esHost = "localhost:1200"
|
|
}
|
|
if !startsWithHTTP(esHost) {
|
|
esHost = "http://" + esHost
|
|
}
|
|
esUser := os.Getenv("ES_USER")
|
|
if esUser == "" {
|
|
esUser = "elastic"
|
|
}
|
|
esPassword := os.Getenv("ES_PASSWORD")
|
|
if esPassword == "" {
|
|
esPassword = "infini_rag_flow"
|
|
}
|
|
|
|
cfg := &server.ElasticsearchConfig{
|
|
Hosts: esHost,
|
|
Username: esUser,
|
|
Password: esPassword,
|
|
}
|
|
docEngine, err = elasticsearch.NewEngine(cfg)
|
|
if err != nil {
|
|
t.Skipf("Could not create Elasticsearch engine: %v (skipping ES subtest)", err)
|
|
return nil, func() {}
|
|
}
|
|
|
|
case engine.EngineInfinity:
|
|
t.Logf("Setting up Infinity engine...")
|
|
infURI := os.Getenv("INFINITY_URI")
|
|
if infURI == "" {
|
|
infURI = "localhost:23817"
|
|
}
|
|
|
|
// First check if Infinity is reachable quickly without waiting 120s!
|
|
if !isPortOpen("localhost", 23817) {
|
|
t.Skipf("Infinity not running at %s (skipping Infinity subtest)", infURI)
|
|
return nil, func() {}
|
|
}
|
|
|
|
cfg := &server.InfinityConfig{
|
|
URI: infURI,
|
|
DBName: "ragflow_e2e_test",
|
|
PostgresPort: 5432,
|
|
}
|
|
docEngine, err = infinity.NewEngine(cfg)
|
|
if err != nil {
|
|
t.Skipf("Could not create Infinity engine: %v (skipping Infinity subtest)", err)
|
|
return nil, func() {}
|
|
}
|
|
|
|
default:
|
|
t.Fatalf("Unsupported engine type: %v", engineType)
|
|
}
|
|
|
|
// Create unique base name for the test
|
|
baseName := fmt.Sprintf("ragflow_%s", tenantID)
|
|
|
|
// Cleanup first (if exists)
|
|
ctx := context.Background()
|
|
_ = docEngine.DropChunkStore(ctx, baseName, datasetID)
|
|
|
|
// Create chunk store (note: vec dimension = 2 because our test chunks have q_2_vec)
|
|
if err := docEngine.CreateChunkStore(ctx, baseName, datasetID, 2, "naive"); err != nil {
|
|
// If create failed, maybe it exists; try dropping and recreating
|
|
_ = docEngine.DropChunkStore(ctx, baseName, datasetID)
|
|
if err := docEngine.CreateChunkStore(ctx, baseName, datasetID, 2, "naive"); err != nil {
|
|
_ = docEngine.Close()
|
|
t.Fatalf("Could not create chunk store: %v", err)
|
|
}
|
|
}
|
|
|
|
// Return cleanup function
|
|
cleanup := func() {
|
|
_ = docEngine.DropChunkStore(ctx, baseName, datasetID)
|
|
_ = docEngine.Close()
|
|
}
|
|
|
|
return docEngine, cleanup
|
|
}
|
|
|
|
func startsWithHTTP(s string) bool {
|
|
return len(s) >= 4 && (s[:4] == "http" || s[:5] == "https")
|
|
}
|
|
|
|
func toLowerSnakeCase(s string) string {
|
|
result := make([]rune, 0, len(s))
|
|
for _, r := range s {
|
|
if r >= 'A' && r <= 'Z' {
|
|
result = append(result, r-'A'+'a')
|
|
} else if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') {
|
|
result = append(result, r)
|
|
} else {
|
|
result = append(result, '_')
|
|
}
|
|
}
|
|
return string(result)
|
|
}
|
|
|
|
func isPortOpen(host string, port int) bool {
|
|
addr := fmt.Sprintf("%s:%d", host, port)
|
|
conn, err := net.DialTimeout("tcp", addr, 1*time.Second)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
defer conn.Close()
|
|
return true
|
|
}
|
|
|
|
// =============================================================================
|
|
// Main E2E Test with Subtests for Elasticsearch and Infinity
|
|
// =============================================================================
|
|
|
|
func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) {
|
|
testCases := []struct {
|
|
name string
|
|
engineType engine.EngineType
|
|
}{
|
|
{
|
|
name: "Elasticsearch",
|
|
engineType: engine.EngineElasticsearch,
|
|
},
|
|
{
|
|
name: "Infinity",
|
|
engineType: engine.EngineInfinity,
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Logf("Running Dataflow E2E test with engine: %s", tc.name)
|
|
|
|
// Setup test database
|
|
db := testutil.SetupTestDB(t)
|
|
cleanupDB := testutil.ReplaceDBForTest(t, db)
|
|
defer cleanupDB()
|
|
|
|
// Seed test data (lowercase for ES index compatibility)
|
|
lowerName := toLowerSnakeCase(tc.name)
|
|
tenantID, kbID, _, taskID := testutil.SeedTestData(t, db,
|
|
testutil.WithTenantID(fmt.Sprintf("tenant-e2e-%s", lowerName)),
|
|
testutil.WithKBID(fmt.Sprintf("kb-e2e-%s", lowerName)),
|
|
testutil.WithDocID(fmt.Sprintf("doc-e2e-%s", lowerName)),
|
|
testutil.WithTaskID(fmt.Sprintf("task-e2e-%s", lowerName)),
|
|
testutil.WithPipelineID(fmt.Sprintf("pipeline-e2e-%s", lowerName)),
|
|
testutil.WithDocName(fmt.Sprintf("e2e-test-%s.pdf", lowerName)),
|
|
)
|
|
|
|
// Setup DocEngine for this test
|
|
docEngine, cleanupEngine := setupTestDocEngine(t, tc.engineType, tenantID, kbID)
|
|
defer cleanupEngine()
|
|
|
|
// Update task to have dataflow task type
|
|
var task entity.Task
|
|
if err := db.Where("id = ?", taskID).First(&task).Error; err != nil {
|
|
t.Fatalf("Failed to get task: %v", err)
|
|
}
|
|
task.TaskType = "dataflow"
|
|
if err := db.Save(&task).Error; err != nil {
|
|
t.Fatalf("Failed to update task: %v", err)
|
|
}
|
|
|
|
// Load task context
|
|
taskCtx, err := LoadTaskContext(taskID)
|
|
if err != nil {
|
|
t.Fatalf("LoadTaskContext failed: %v", err)
|
|
}
|
|
|
|
// Track what was called
|
|
var (
|
|
loadDSLCalled bool
|
|
runPipelineCalled bool
|
|
insertChunksCalled bool
|
|
)
|
|
var capturedChunks [][]map[string]any
|
|
|
|
// Create TaskHandler with mocked DataflowService factory
|
|
handler := NewTaskHandler(taskCtx)
|
|
handler.WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) {
|
|
svc := mustNewDataflowService(t, ctx, dataflowID, 0, 0)
|
|
|
|
// Mock loadDSLFunc
|
|
svc.WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) {
|
|
loadDSLCalled = true
|
|
return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, dataflowID, nil
|
|
})
|
|
|
|
// Mock runPipelineFunc - returns test chunks (with vectors to skip embedding)
|
|
svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) {
|
|
runPipelineCalled = true
|
|
return map[string]any{
|
|
"chunks": []map[string]any{
|
|
{
|
|
"text": fmt.Sprintf("Hello world from E2E test with %s", tc.name),
|
|
"id": fmt.Sprintf("chunk-e2e-%s-1", tc.name),
|
|
"q_2_vec": []float64{0.1, 0.2}, // Pre-vectorized to skip embedding
|
|
},
|
|
{
|
|
"text": fmt.Sprintf("Second chunk from E2E test with %s", tc.name),
|
|
"id": fmt.Sprintf("chunk-e2e-%s-2", tc.name),
|
|
"q_2_vec": []float64{0.3, 0.4}, // Pre-vectorized to skip embedding
|
|
},
|
|
},
|
|
EmbeddingTokenConsumptionKey: 100,
|
|
}, dsl, nil
|
|
})
|
|
|
|
// Use the injected DocEngine for insertChunks!
|
|
svc.WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) {
|
|
insertChunksCalled = true
|
|
t.Logf("DocEngine InsertChunks called! baseName=%s datasetID=%s len(chunks)=%d", baseName, datasetID, len(chunks))
|
|
ids, err := docEngine.InsertChunks(ctx, chunks, baseName, datasetID)
|
|
if err != nil {
|
|
t.Logf("WARNING: InsertChunks err=%v", err)
|
|
}
|
|
capturedChunks = append(capturedChunks, chunks)
|
|
return ids, err
|
|
})
|
|
|
|
svc.WithChunkCounter(&stubChunkCounter{})
|
|
return svc, nil
|
|
})
|
|
|
|
// Also set progress callback
|
|
var progressEvents []struct {
|
|
prog float64
|
|
msg string
|
|
}
|
|
taskCtx.ProgressFunc = func(prog float64, msg string) {
|
|
t.Logf("PROGRESS: %.2f %s", prog, msg)
|
|
progressEvents = append(progressEvents, struct {
|
|
prog float64
|
|
msg string
|
|
}{prog, msg})
|
|
}
|
|
|
|
// Execute the task handler!
|
|
t.Logf("Calling TaskHandler.Handle()...")
|
|
err = handler.Handle()
|
|
if err != nil {
|
|
t.Fatalf("TaskHandler.Handle failed: %v", err)
|
|
}
|
|
t.Logf("TaskHandler.Handle() complete!")
|
|
|
|
// Verify all the expected calls happened
|
|
if !loadDSLCalled {
|
|
t.Fatal("Expected loadDSLFunc to be called")
|
|
}
|
|
if !runPipelineCalled {
|
|
t.Fatal("Expected runPipelineFunc to be called")
|
|
}
|
|
if !insertChunksCalled {
|
|
t.Fatal("Expected insertChunks to be called")
|
|
}
|
|
|
|
// Verify chunks were captured
|
|
totalChunks := 0
|
|
for _, batch := range capturedChunks {
|
|
totalChunks += len(batch)
|
|
}
|
|
if totalChunks != 2 {
|
|
t.Fatalf("Expected total 2 chunks, got %d", totalChunks)
|
|
}
|
|
|
|
// Now verify we can read the chunks back!
|
|
baseName := fmt.Sprintf("ragflow_%s", tenantID)
|
|
t.Logf("Reading chunks back from %s: baseName=%s datasetID=%s", tc.name, baseName, kbID)
|
|
|
|
// Refresh (wait for ES to index or Infinity to commit)
|
|
time.Sleep(300 * time.Millisecond)
|
|
|
|
for _, batch := range capturedChunks {
|
|
for _, chunk := range batch {
|
|
chunkID, ok := chunk["id"].(string)
|
|
if ok && chunkID != "" {
|
|
t.Logf("Trying to get chunk: %s", chunkID)
|
|
readBack, err := docEngine.GetChunk(context.Background(), baseName, chunkID, []string{kbID})
|
|
if err != nil {
|
|
t.Logf("WARNING: Failed to get chunk %q: %v", chunkID, err)
|
|
} else if readBack == nil {
|
|
t.Logf("WARNING: chunk %q not found (may not have been indexed yet)", chunkID)
|
|
} else {
|
|
t.Logf("SUCCESS: Read back chunk %q!", chunkID)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Verify progress reported
|
|
if len(progressEvents) == 0 {
|
|
t.Fatal("Expected at least one progress event")
|
|
}
|
|
foundDone := false
|
|
for _, ev := range progressEvents {
|
|
if ev.prog == 1.0 {
|
|
foundDone = true
|
|
}
|
|
}
|
|
if !foundDone {
|
|
t.Fatal("Expected progress to reach 1.0")
|
|
}
|
|
|
|
// Verify final task status can be updated to success
|
|
ingestionTaskDAO := dao.NewIngestionTaskDAO()
|
|
if err := ingestionTaskDAO.UpdateStatus(taskID, "success"); err != nil {
|
|
t.Fatalf("UpdateStatus failed: %v", err)
|
|
}
|
|
|
|
finalTask, err := ingestionTaskDAO.GetByID(taskID)
|
|
if err != nil {
|
|
t.Fatalf("GetByID failed: %v", err)
|
|
}
|
|
if finalTask.Status != "success" {
|
|
t.Errorf("Final task status = %q, want %q", finalTask.Status, "success")
|
|
}
|
|
|
|
t.Logf("SUCCESS: Dataflow E2E test passed with %s engine!", tc.name)
|
|
})
|
|
}
|
|
}
|