build-flow.api/tests/BuildFlow.Rag.Tests/PreviewDocumentChunksQueryHandlerTests.cs

using BuildFlow.Application.Document.Handlers;
using BuildFlow.Application.Document.Queries;
using BuildFlow.Domain.Entities.Document;
using BuildFlow.Domain.Entities.KnowledgeBase;
using BuildFlow.Domain.Enums;
using BuildFlow.Domain.Interfaces;
using BuildFlow.Infrastructure.AI;
using BuildFlow.Infrastructure.Persistence;
using BuildFlow.Rag.Tests.Stub;
using File.Application.Features.Files.DTOs;
using File.Application.Features.Files.Queries;
using MediatR;
using Microsoft.EntityFrameworkCore;
using Moq;
using Xunit;

namespace BuildFlow.Rag.Tests;

/// <summary>
/// 文档切片预览测试。预览只提取并分块，不生成 embedding，也不写入正式分块表。
/// </summary>
public class PreviewDocumentChunksQueryHandlerTests
{
    [Fact]
    public async Task HandleShouldReturnDryRunChunksWithoutPersistingDocumentChunks()
    {
        var options = new DbContextOptionsBuilder<RagDbContext>()
            .UseInMemoryDatabase($"rag-preview-{Guid.NewGuid():N}")
            .Options;
        await using var db = new RagDbContext(options);
        var kbId = Guid.NewGuid();
        var docId = Guid.NewGuid();
        var filePath = Path.Combine(Path.GetTempPath(), $"preview-{Guid.NewGuid():N}.md");
        await System.IO.File.WriteAllTextAsync(filePath, "# 第一章\n\nBuildFlow 文档切片预览。\n\n# 第二章\n\n上传后先看切片，再手动处理向量。");

        try
        {
            await db.KnowledgeBases.AddAsync(new KnowledgeBaseEntity
            {
                Id = kbId,
                Name = "预览知识库",
                ChunkSize = 40,
                ChunkOverlap = 5,
                ChunkingMode = ChunkingModeEnum.Heading,
                EmbeddingModel = "text-embedding-embeddinggemma-300m",
                Status = KnowledgeBaseStatusEnum.Active,
                CreatedBy = "test",
                UpdatedBy = "test",
                OperatorIP = "127.0.0.1",
            });
            await db.Documents.AddAsync(new DocumentEntity
            {
                Id = docId,
                KnowledgeBaseId = kbId,
                Title = "preview.md",
                FileName = "preview.md",
                FilePath = filePath,
                FileSize = 64,
                ContentType = "text/markdown",
                Status = DocumentStatusEnum.Pending,
                CreatedBy = "test",
                UpdatedBy = "test",
                OperatorIP = "127.0.0.1",
            });
            await db.SaveChangesAsync();

            var mediator = new Mock<IMediator>();
            var handler = new PreviewDocumentChunksQueryHandler(
                db,
                new StubDocumentExtractorFactory(),
                new TextChunker(),
                mediator.Object);

            var result = await handler.Handle(new PreviewDocumentChunksQuery(docId), CancellationToken.None);

            Assert.Equal(2, result.Count);
            Assert.All(result, chunk => Assert.False(chunk.HasEmbedding));
            Assert.Contains("第一章", result[0].Content);
            Assert.Contains("第二章", result[1].Content);
            Assert.Equal(0, await db.DocumentChunks.CountAsync());
            mediator.Verify(m => m.Send(
                It.IsAny<DownloadFileQuery>(),
                It.IsAny<CancellationToken>()), Times.Never);
        }
        finally
        {
            System.IO.File.Delete(filePath);
        }
    }

    [Fact]
    public async Task HandleRemoteDocumentShouldDownloadThroughFileQueryAndDisposeStream()
    {
        var options = new DbContextOptionsBuilder<RagDbContext>()
            .UseInMemoryDatabase($"rag-preview-remote-{Guid.NewGuid():N}")
            .Options;
        await using var db = new RagDbContext(options);
        var kbId = Guid.NewGuid();
        var docId = Guid.NewGuid();
        await db.KnowledgeBases.AddAsync(new KnowledgeBaseEntity
        {
            Id = kbId,
            Name = "远端预览知识库",
            ChunkSize = 40,
            ChunkOverlap = 5,
            ChunkingMode = ChunkingModeEnum.Heading,
            Status = KnowledgeBaseStatusEnum.Active,
            CreatedBy = "test",
            UpdatedBy = "test",
            OperatorIP = "127.0.0.1",
        });
        await db.Documents.AddAsync(new DocumentEntity
        {
            Id = docId,
            KnowledgeBaseId = kbId,
            Title = "remote.md",
            FileName = "remote.md",
            FilePath = "rag-documents/docs/kb/remote.md",
            FileSize = 32,
            ContentType = "text/markdown",
            Status = DocumentStatusEnum.Pending,
            CreatedBy = "test",
            UpdatedBy = "test",
            OperatorIP = "127.0.0.1",
        });
        await db.SaveChangesAsync();
        var stream = new TrackingMemoryStream("# 第一章\n\n远端文档内容。"u8.ToArray());
        var mediator = new Mock<IMediator>();
        mediator.Setup(m => m.Send(
                It.Is<DownloadFileQuery>(query => query.ObjectKey == "docs/kb/remote.md"),
                It.IsAny<CancellationToken>()))
            .ReturnsAsync(new DownloadFileResult(stream, "text/markdown", "remote.md"));
        var handler = new PreviewDocumentChunksQueryHandler(
            db,
            new StubDocumentExtractorFactory(),
            new TextChunker(),
            mediator.Object);

        var result = await handler.Handle(new PreviewDocumentChunksQuery(docId), CancellationToken.None);

        Assert.NotEmpty(result);
        Assert.True(stream.IsDisposed);
        mediator.Verify(m => m.Send(
            It.Is<DownloadFileQuery>(query => query.ObjectKey == "docs/kb/remote.md"),
            It.IsAny<CancellationToken>()), Times.Once);
    }

    private sealed class StubDocumentExtractorFactory : IDocumentExtractorFactory
    {
        public IDocumentExtractor GetExtractor(string fileNameOrExtension) => new StubDocumentExtractor();
    }

    private sealed class StubDocumentExtractor : IDocumentExtractor
    {
        public IReadOnlyCollection<string> SupportedExtensions { get; init; } = ["md"];

        public async Task<ExtractedDocument> ExtractAsync(string filePath, CancellationToken ct = default)
        {
            return new ExtractedDocument
            {
                Content = await System.IO.File.ReadAllTextAsync(filePath, ct),
                Title = Path.GetFileName(filePath),
            };
        }
    }
}