internal/ingestion/pipeline/media_real_integration_test.go

Rama main · réplica de lectura del repositorio original.

//go:build integration

package pipeline

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/cespare/xxhash/v2"
	"gopkg.in/yaml.v3"
	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/engine/kvrocks"
	"ragflow/internal/entity"
	component "ragflow/internal/ingestion/component"
	"ragflow/internal/ingestion/component/schema"
	"ragflow/internal/ingestion/mediacache"
	"ragflow/internal/parser/media"
	"ragflow/internal/parser/parser"
	"ragflow/internal/server"
	"ragflow/internal/server/config"
	"ragflow/internal/storage"
	"ragflow/internal/tokenizer"
)

// runMediaRealServices exercises the normal canvas with real MySQL, Kvrocks
// and MinIO, checking parser references, uploaded bytes and cache lifecycle.
func runMediaRealServices(t *testing.T, format string, source, png []byte) {
	if common.GetEnv(common.EnvRAGFlowDictPath) == "" {
		const dictPath = "/usr/share/infinity/resource"
		if _, err := os.Stat(filepath.Join(dictPath, "rag", "huqie.txt")); err != nil {
			t.Fatalf("tokenizer resource: %v", err)
		}
		t.Setenv(common.EnvRAGFlowDictPath, dictPath)
	}
	if err := tokenizer.Init(&tokenizer.PoolConfig{DictPath: common.GetEnv(common.EnvRAGFlowDictPath), MinSize: 1, MaxSize: 2, IdleTimeout: 30 * time.Second, AcquireTimeout: 5 * time.Second}); err != nil {
		t.Fatalf("tokenizer initialization: %v", err)
	}
	cfg := loadMediaIntegrationConfig(t)
	db := mustOpenRealMySQL(t, cfg)
	sqlDB, err := db.DB()
	if err != nil {
		t.Fatal(err)
	}
	sqlDB.SetMaxOpenConns(1)
	if err := db.Exec("SET SESSION sql_mode = ''").Error; err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = sqlDB.Close() })
	if err := db.AutoMigrate(&entity.Tenant{}, &entity.UserTenant{}, &entity.Knowledgebase{}, &entity.Document{}, &entity.File{}, &entity.File2Document{}); err != nil {
		t.Fatal(err)
	}
	if err := kvrocks.Init(t.Context()); err != nil {
		t.Fatal(err)
	}
	kv := kvrocks.Get()
	stg, err := storage.NewMinioStorage(cfg.GetMinioConfig())
	if err != nil {
		t.Fatal(err)
	}
	readyCtx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
	defer cancel()
	for {
		_, err := stg.BucketExistsWithError(readyCtx, "ragflow-html-media-probe")
		if err == nil {
			break
		}
		select {
		case <-readyCtx.Done():
			t.Fatalf("MinIO S3 not ready: %v", err)
		case <-time.After(250 * time.Millisecond):
		}
	}
	originalDB := dao.DB
	originalStorage := storage.GetStorageFactory().GetStorage()
	originalResolver := component.ResolveDocumentStorageOverride
	dao.DB = db
	storage.GetStorageFactory().SetStorage(stg)
	component.ResolveDocumentStorageOverride = nil
	t.Cleanup(func() {
		dao.DB = originalDB
		storage.GetStorageFactory().SetStorage(originalStorage)
		component.ResolveDocumentStorageOverride = originalResolver
	})
	template, err := os.ReadFile(filepath.Join(repoRootFromPipelineTest(t), "internal/ingestion/pipeline/template/ingestion_pipeline_general.json"))
	if err != nil {
		t.Fatal(err)
	}
	template = disableTokenizerEmbeddingForTemplate(t, template)
	suffix := fmt.Sprintf("%d", time.Now().UnixNano())
	tenantID, kbID := limit32("it_tenant_"+suffix), s3SafeBucketName("it-kb-"+suffix)
	docID, fileID := limit32("it_doc_"+suffix), limit32("it_file_"+suffix)
	taskID := "it-" + format + "-" + suffix
	bucket := s3SafeBucketName(kbID)
	name := "image." + format
	objectPath := "integration/" + taskID + "/" + name
	mustSeedRealPipelineDocument(t, db, stg, tenantID, kbID, docID, fileID, bucket, objectPath, name, string(source))
	t.Cleanup(func() {
		_ = mediacache.PurgeTask(context.Background(), kv, taskID)
		cleanupRealPipelineDocument(db, stg, tenantID, kbID, docID, fileID, bucket, objectPath)
	})
	if err := db.Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]any{"type": format, "suffix": "." + format}).Error; err != nil {
		t.Fatal(err)
	}
	if err := db.Model(&entity.File{}).Where("id = ?", fileID).Update("type", format).Error; err != nil {
		t.Fatal(err)
	}
	pipe, err := NewPipelineFromDSL(template, taskID, WithDocumentID(docID), WithKV(kv))
	if err != nil {
		t.Fatal(err)
	}
	out, err := pipe.Run(t.Context(), map[string]any{"doc_id": docID, "kb_id": kbID, "tenant_id": tenantID}, nil)
	if err != nil {
		t.Fatal(err)
	}
	state := stateFromRunOutput(t, out)
	docs, ok := state["Parser:HipSignsRhyme"]["json"].([]parser.DocItem)
	if !ok {
		t.Fatalf("parser output = %T", state["Parser:HipSignsRhyme"]["json"])
	}
	var ref *entity.MediaRef
	for _, item := range docs {
		m := item.ToMap()
		if m["media_ref"] != nil {
			ref = entity.MediaRefFromAny(m["media_ref"])
			if m["image"] != nil {
				t.Fatal("parser retained inline image")
			}
		}
	}
	if ref == nil || ref.Format != format {
		t.Fatal("parser did not emit media_ref")
	}
	cached, err := mediacache.New(kv, taskID).Get(t.Context(), *ref)
	if err != nil || !bytes.Equal(cached, png) {
		t.Fatalf("cached bytes differ: %v", err)
	}
	chunks, ok := schema.ChunkCarrierToMaps(state["GeneralChunker:SixApplesFall"]["chunks"])
	if !ok {
		t.Fatal("missing chunker output")
	}
	foundImage := false
	for _, ck := range chunks {
		id, _ := ck["id"].(string)
		imgID, _ := ck["img_id"].(string)
		if imgID == "" {
			continue
		}
		if imgID != kbID+"-"+id || ck["media_ref"] != nil {
			t.Fatalf("invalid image chunk: %v", ck)
		}
		got, err := stg.Get(t.Context(), kbID, id)
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { _ = stg.Remove(context.Background(), kbID, id) })
		if !bytes.Equal(got, png) {
			t.Fatal("stored image does not match source")
		}
		foundImage = true
	}
	if !foundImage {
		t.Fatalf("chunker produced no uploaded image: chunks=%#v state=%#v", chunks, state["GeneralChunker:SixApplesFall"])
	}
	key := fmt.Sprintf("mc:media:%x:%s", xxhash.Sum64String(taskID), ref.SHA256)
	ttl, err := kv.GetClient().TTL(t.Context(), key).Result()
	if err != nil || ttl <= 23*time.Hour || ttl > mediacache.TTL {
		t.Fatalf("media TTL = %v, %v", ttl, err)
	}
	resolveCtx := media.WithCache(t.Context(), mediacache.New(kv, taskID))
	if err := kv.GetClient().Set(t.Context(), key, "corrupt", mediacache.TTL).Err(); err != nil {
		t.Fatal(err)
	}
	if _, _, err := component.ResolveMediaRef(resolveCtx, map[string]any{"media_ref": ref}); err == nil || errors.Is(err, media.ErrMediaNotFound) {
		t.Fatalf("corrupt payload must propagate an integrity error: %v", err)
	}
	if err := kv.GetClient().Set(t.Context(), key, png, mediacache.TTL).Err(); err != nil {
		t.Fatal(err)
	}
	if err := kv.GetClient().Expire(t.Context(), key, time.Second).Err(); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(2 * time.Second)
	for {
		_, err = mediacache.New(kv, taskID).Get(t.Context(), *ref)
		if errors.Is(err, media.ErrMediaNotFound) {
			break
		}
		if err != nil || time.Now().After(deadline) {
			t.Fatalf("media did not expire: %v", err)
		}
		time.Sleep(10 * time.Millisecond)
	}
	if _, _, err := component.ResolveMediaRef(resolveCtx, map[string]any{"media_ref": ref}); !errors.Is(err, media.ErrMediaNotFound) {
		t.Fatalf("expired reference must propagate a cache miss: %v", err)
	}
	retry := mediacache.New(kv, taskID)
	retryRef, err := retry.Put(t.Context(), format, "image/png", png)
	if err != nil {
		t.Fatal(err)
	}
	if got, err := retry.Get(t.Context(), retryRef); err != nil || !bytes.Equal(got, png) {
		t.Fatalf("retry cache bytes: %v", err)
	}
	scopeHash := xxhash.Sum64String(taskID)
	digestKey := fmt.Sprintf("mc:quota:%x:digests", scopeHash)
	stateKey := fmt.Sprintf("mc:quota:%x:state", scopeHash)
	manifestKey := fmt.Sprintf("mc:manifest:%x", scopeHash)
	if member, err := kv.GetClient().SIsMember(t.Context(), digestKey, retryRef.SHA256).Result(); err != nil || !member {
		t.Fatalf("quota digest membership = %v, %v; want true", member, err)
	}
	images, err := kv.GetClient().SCard(t.Context(), digestKey).Result()
	if err != nil || images != 1 {
		t.Fatalf("quota images = %d, %v; want 1", images, err)
	}
	bytesUsed, err := kv.GetClient().HGet(t.Context(), stateKey, "bytes").Result()
	if err != nil || bytesUsed != fmt.Sprint(len(png)) {
		t.Fatalf("quota bytes = %q, %v; want %d", bytesUsed, err, len(png))
	}
	if member, err := kv.GetClient().SIsMember(t.Context(), manifestKey, key).Result(); err != nil || !member {
		t.Fatalf("manifest membership = %v, %v; want true", member, err)
	}
	if err := mediacache.PurgeTask(t.Context(), kv, taskID); err != nil {
		t.Fatal(err)
	}
	if _, err := retry.Get(t.Context(), retryRef); !errors.Is(err, media.ErrMediaNotFound) {
		t.Fatalf("purged cache read: %v", err)
	}
	for _, purgedKey := range []string{key, digestKey, stateKey, manifestKey} {
		if exists, err := kv.GetClient().Exists(t.Context(), purgedKey).Result(); err != nil || exists != 0 {
			t.Fatalf("purged key %s exists=%d, err=%v", purgedKey, exists, err)
		}
	}
}

// loadMediaIntegrationConfig uses the CI service addresses when set.
func loadMediaIntegrationConfig(t *testing.T) *config.Config {
	t.Helper()
	kvAddr := os.Getenv("RAGFLOW_TEST_KVROCKS_ADDR")
	if kvAddr == "" {
		return mustLoadRealIntegrationConfig(t)
	}
	cfg := map[string]any{
		"general": map[string]any{"cache_engine": "kvrocks"},
		"mysql":   map[string]any{"host": os.Getenv("RAGFLOW_MYSQL_HOST"), "port": os.Getenv("RAGFLOW_MYSQL_PORT"), "password": os.Getenv("MYSQL_PASSWORD")},
		"minio":   map[string]any{"host": os.Getenv("RAGFLOW_MINIO_HOST"), "user": os.Getenv("MINIO_USER"), "password": os.Getenv("MINIO_PASSWORD")},
		"kvrocks": map[string]any{"host": kvAddr, "password": os.Getenv("REDIS_PASSWORD"), "db": 1},
	}
	data, err := yaml.Marshal(cfg)
	if err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(t.TempDir(), "service_conf.yaml")
	if err := os.WriteFile(path, data, 0600); err != nil {
		t.Fatal(err)
	}
	if err := common.InitLogger("info", common.FileOutput{}, ""); err != nil {
		t.Fatal(err)
	}
	if err := server.Init(path); err != nil {
		t.Fatal(err)
	}
	return server.GetConfig()
}