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

using BuildFlow.Application.Chat.Commands;
using BuildFlow.Application.Chat.DTOs;
using BuildFlow.Application.Chat.Handlers;
using BuildFlow.Domain.Entities.ChatMessage;
using BuildFlow.Domain.Entities.Conversation;
using BuildFlow.Domain.Enums;
using BuildFlow.Infrastructure.Persistence;
using BuildFlow.Rag.Tests.Stub;
using BuildFlow.Shared.Domain.Errors;
using BuildFlow.Shared.Domain.Exceptions;
using BuildFlow.Shared.Infrastructure;
using FluentAssertions;
using Microsoft.EntityFrameworkCore;
using Xunit;

namespace BuildFlow.Rag.Tests;

/// <summary>
/// RAG 流式发送和编辑重生成的对象所有权与事件契约测试。
/// </summary>
public class ChatStreamOwnershipTests
{
    [Fact]
    public async Task StreamMessageForOwnerShouldYieldChunksThenDoneAndPersistMessages()
    {
        var userContext = new TestCurrentUserContext();
        await using var db = CreateDb(userContext);
        var conversationId = await SeedConversationAsync(
            db,
            Guid.Parse(userContext.UserId));
        var cache = new StubChatMessageCache();
        var handler = new StreamMessageCommandHandler(
            db,
            new StubAIChatAgent { StreamTokens = ["你", "好", "！"] },
            cache,
            userContext);

        var events = await CollectAsync(handler.Handle(
            new StreamMessageCommand(conversationId, "你好"),
            CancellationToken.None));

        events.Take(3).Select(x => x.Content).Should().Equal("你", "好", "！");
        events.Should().HaveCount(4);
        events[^1].MessageId.Should().NotBeNull();
        events[^1].MessageId!.Value.Should().NotBeEmpty();
        db.ChatMessages.Count(x => x.Role == ChatRoleEnum.User).Should().Be(1);
        db.ChatMessages.Count(x => x.Role == ChatRoleEnum.Assistant).Should().Be(1);
        cache.Count(conversationId).Should().Be(2);
    }

    [Fact]
    public async Task StreamMessageForNonOwnerShouldThrowNotFoundBeforeWriting()
    {
        var userContext = new TestCurrentUserContext();
        await using var db = CreateDb(userContext);
        var conversationId = await SeedConversationAsync(db, Guid.NewGuid());
        var cache = new StubChatMessageCache();
        var handler = new StreamMessageCommandHandler(
            db,
            new StubAIChatAgent(),
            cache,
            userContext);

        var act = () => CollectAsync(handler.Handle(
            new StreamMessageCommand(conversationId, "越权消息"),
            CancellationToken.None));

        await act.Should().ThrowAsync<NotFoundException>()
            .Where(x => x.Code == ErrorCode.NotFound && x.Message == "会话不存在");
        db.ChatMessages.Should().BeEmpty();
        cache.Count(conversationId).Should().Be(0);
    }

    [Fact]
    public async Task EditStreamForNonOwnerShouldThrowNotFoundBeforeEditing()
    {
        var userContext = new TestCurrentUserContext();
        await using var db = CreateDb(userContext);
        var conversationId = await SeedConversationAsync(db, Guid.NewGuid());
        var target = new ChatMessageEntity
        {
            ConversationId = conversationId,
            Role = ChatRoleEnum.User,
            Content = "原始问题",
        };
        var later = new ChatMessageEntity
        {
            ConversationId = conversationId,
            Role = ChatRoleEnum.Assistant,
            Content = "原始回答",
        };
        db.ChatMessages.AddRange(target, later);
        await db.SaveChangesAsync();
        var handler = new EditMessageStreamCommandHandler(
            db,
            new StubAIChatAgent(),
            new StubChatMessageCache(),
            userContext);

        var act = () => CollectAsync(handler.Handle(
            new EditMessageStreamCommand(conversationId, target.Id, "越权编辑"),
            CancellationToken.None));

        await act.Should().ThrowAsync<NotFoundException>()
            .Where(x => x.Code == ErrorCode.NotFound && x.Message == "会话不存在");
        target.Content.Should().Be("原始问题");
        db.ChatMessages.Should().HaveCount(2);
    }

    [Fact]
    public async Task EditStreamForOwnerShouldPreserveChunkAndDoneEventContract()
    {
        var userContext = new TestCurrentUserContext();
        await using var db = CreateDb(userContext);
        var conversationId = await SeedConversationAsync(
            db,
            Guid.Parse(userContext.UserId));
        var target = new ChatMessageEntity
        {
            ConversationId = conversationId,
            Role = ChatRoleEnum.User,
            Content = "原始问题",
        };
        db.ChatMessages.Add(target);
        await db.SaveChangesAsync();
        await Task.Delay(5);
        db.ChatMessages.Add(new ChatMessageEntity
        {
            ConversationId = conversationId,
            Role = ChatRoleEnum.Assistant,
            Content = "旧回答",
        });
        await db.SaveChangesAsync();
        var cache = new StubChatMessageCache();
        var handler = new EditMessageStreamCommandHandler(
            db,
            new StubAIChatAgent { StreamTokens = ["新", "回", "答"] },
            cache,
            userContext);

        var events = await CollectAsync(handler.Handle(
            new EditMessageStreamCommand(conversationId, target.Id, "新问题"),
            CancellationToken.None));

        events.Take(3).Select(x => x.Content).Should().Equal("新", "回", "答");
        events.Should().HaveCount(4);
        events[^1].MessageId.Should().NotBeNull();
        events[^1].MessageId!.Value.Should().NotBeEmpty();
        target.Content.Should().Be("新问题");
        db.ChatMessages.Should().HaveCount(2);
        db.ChatMessages.Single(x => x.Role == ChatRoleEnum.Assistant)
            .Content.Should().Be("新回答");
        cache.Count(conversationId).Should().Be(2);
    }

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

    private static async Task<Guid> SeedConversationAsync(RagDbContext db, Guid userId)
    {
        var conversation = new ConversationEntity
        {
            Title = "测试会话",
            UserId = userId,
        };
        db.Conversations.Add(conversation);
        await db.SaveChangesAsync();
        return conversation.Id;
    }

    private static async Task<List<T>> CollectAsync<T>(IAsyncEnumerable<T> source)
    {
        var result = new List<T>();
        await foreach (var item in source) result.Add(item);
        return result;
    }
}