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

using BuildFlow.Application.RagQA.DTOs;
using BuildFlow.Application.RagQA.Handlers;
using BuildFlow.Application.RagQA.Queries;
using BuildFlow.Application.RagQA.Validators;
using BuildFlow.Domain.Entities.KnowledgeBase;
using BuildFlow.Domain.Interfaces;
using BuildFlow.Infrastructure.Persistence;
using BuildFlow.Rag.Tests.Stub;
using BuildFlow.Shared.Domain.Errors;
using BuildFlow.Shared.Domain.Exceptions;
using FluentAssertions;
using Microsoft.EntityFrameworkCore;
using Xunit;

namespace BuildFlow.Rag.Tests;

/// <summary>
/// RAG 流式查询应用层测试，固定来源、内容、结束事件顺序和错误语义。
/// </summary>
public class RagStreamQueryHandlerTests
{
    [Fact]
    public async Task RetrievedSourcesAreFollowedByContentChunksAndDone()
    {
        await using var db = CreateDb();
        var knowledgeBase = await SeedKnowledgeBaseAsync(db);
        var documentId = Guid.NewGuid();
        var longContext = new string('a', 220);
        var retrieval = new StubRagRetrievalService
        {
            Hits =
            [
                new RetrievalHit(
                    documentId,
                    "架构说明",
                    Guid.NewGuid(),
                    longContext,
                    0.87654,
                    longContext),
            ],
        };
        var handler = new RAGStreamQueryHandler(
            db,
            retrieval,
            new StubAIChatAgent { StreamTokens = ["你", "好"] });

        var events = await CollectAsync(handler.Handle(
            new RAGStreamQuery(knowledgeBase.Id, "系统如何设计？"),
            CancellationToken.None));

        events.Should().HaveCount(4);
        events[0].Sources.Should().ContainSingle();
        var source = events[0].Sources!.Single();
        source.DocumentId.Should().Be(documentId);
        source.Content.Should().HaveLength(203);
        source.Content.Should().EndWith("...");
        source.Similarity.Should().Be(0.8765);
        events[1].Content.Should().Be("你");
        events[2].Content.Should().Be("好");
        events[3].IsDone.Should().BeTrue();
    }

    [Fact]
    public async Task EmptyRetrievalReturnsFallbackContentAndDoneWithoutCallingChatAgent()
    {
        await using var db = CreateDb();
        var knowledgeBase = await SeedKnowledgeBaseAsync(db);
        var chatAgent = new RecordingChatAgent();
        var handler = new RAGStreamQueryHandler(
            db,
            new StubRagRetrievalService { Hits = [] },
            chatAgent);

        var events = await CollectAsync(handler.Handle(
            new RAGStreamQuery(knowledgeBase.Id, "没有答案的问题"),
            CancellationToken.None));

        events.Should().HaveCount(2);
        events[0].Content.Should().Be("未在知识库中找到相关内容。");
        events[1].IsDone.Should().BeTrue();
        chatAgent.StreamCallCount.Should().Be(0);
    }

    [Fact]
    public async Task MissingKnowledgeBaseReturnsNotFoundCode()
    {
        await using var db = CreateDb();
        var handler = new RAGStreamQueryHandler(
            db,
            new StubRagRetrievalService(),
            new StubAIChatAgent());

        var act = () => CollectAsync(handler.Handle(
            new RAGStreamQuery(Guid.NewGuid(), "问题"),
            CancellationToken.None));

        var exception = await act.Should().ThrowAsync<NotFoundException>();
        exception.Which.Code.Should().Be(ErrorCode.NotFound);
    }

    [Fact]
    public async Task CancellationIsPropagatedWhileStreaming()
    {
        await using var db = CreateDb();
        var knowledgeBase = await SeedKnowledgeBaseAsync(db);
        var retrieval = new StubRagRetrievalService
        {
            Hits =
            [
                new RetrievalHit(
                    Guid.NewGuid(),
                    "文档",
                    Guid.NewGuid(),
                    "上下文",
                    0.9,
                    "上下文"),
            ],
        };
        var handler = new RAGStreamQueryHandler(
            db,
            retrieval,
            new StubAIChatAgent { StreamTokens = ["第一块", "第二块"] });
        using var cts = new CancellationTokenSource();
        cts.Cancel();

        var act = () => CollectAsync(handler.Handle(
            new RAGStreamQuery(knowledgeBase.Id, "问题"),
            cts.Token));

        await act.Should().ThrowAsync<OperationCanceledException>();
    }

    [Fact]
    public void StreamQueryValidatorAppliesKnowledgeBaseAndQuestionLimits()
    {
        var validator = new RAGStreamQueryValidator();

        var result = validator.Validate(new RAGStreamQuery(Guid.Empty, new string('x', 1001)));

        result.IsValid.Should().BeFalse();
        result.Errors.Should().Contain(error => error.ErrorMessage == "知识库ID不能为空");
        result.Errors.Should().Contain(error => error.ErrorMessage == "问题不能超过1000个字符");
    }

    private static RagDbContext CreateDb()
    {
        var options = new DbContextOptionsBuilder<RagDbContext>()
            .UseInMemoryDatabase($"rag-stream-{Guid.NewGuid():N}")
            .Options;
        return new RagDbContext(options);
    }

    private static async Task<KnowledgeBaseEntity> SeedKnowledgeBaseAsync(RagDbContext db)
    {
        var knowledgeBase = new KnowledgeBaseEntity
        {
            Name = "流式知识库",
            CreatedBy = "test",
            UpdatedBy = "test",
            OperatorIP = "127.0.0.1",
        };
        db.KnowledgeBases.Add(knowledgeBase);
        await db.SaveChangesAsync();
        return knowledgeBase;
    }

    private static async Task<List<RAGStreamEvent>> CollectAsync(
        IAsyncEnumerable<RAGStreamEvent> stream)
    {
        var events = new List<RAGStreamEvent>();
        await foreach (var streamEvent in stream)
        {
            events.Add(streamEvent);
        }

        return events;
    }

    private sealed class RecordingChatAgent : IAIChatAgent
    {
        public int StreamCallCount { get; private set; }

        public Task<string> RunAsync(
            IReadOnlyList<Microsoft.Extensions.AI.ChatMessage> messages,
            CancellationToken ct = default) => Task.FromResult(string.Empty);

        public async IAsyncEnumerable<string> RunStreamingAsync(
            IReadOnlyList<Microsoft.Extensions.AI.ChatMessage> messages,
            [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken ct = default)
        {
            StreamCallCount++;
            await Task.Yield();
            yield return "unexpected";
        }
    }
}