64 lines
2.4 KiB
C#
64 lines
2.4 KiB
C#
using System.Text.Json;
|
|
using BuildFlow.Application.RagQA.DTOs;
|
|
using BuildFlow.Application.RagQA.Queries;
|
|
using BuildFlow.Host.Endpoints.Rag.RAG;
|
|
using FastEndpoints;
|
|
using FluentAssertions;
|
|
using MediatR;
|
|
using NSubstitute;
|
|
using Xunit;
|
|
|
|
namespace BuildFlow.Host.Tests;
|
|
|
|
/// <summary>
|
|
/// RAG 流式 Endpoint 测试，固定既有 Server-Sent Events 协议。
|
|
/// </summary>
|
|
public class RagStreamEndpointTests
|
|
{
|
|
[Fact]
|
|
public async Task EndpointEncodesSourcesContentAndDoneUsingExistingSseProtocol()
|
|
{
|
|
var documentId = Guid.NewGuid();
|
|
var sources = new List<SourceChunk>
|
|
{
|
|
new(documentId, "架构说明", "来源片段", 0.9),
|
|
};
|
|
var mediator = Substitute.For<IMediator>();
|
|
mediator.CreateStream(Arg.Any<RAGStreamQuery>(), Arg.Any<CancellationToken>())
|
|
.Returns(CreateEvents(sources));
|
|
await using var body = new MemoryStream();
|
|
var endpoint = Factory.Create<RAGStreamEndpoint>(
|
|
context => context.Response.Body = body,
|
|
mediator);
|
|
var request = new RAGQueryRequest(Guid.NewGuid(), "问题");
|
|
|
|
await endpoint.HandleAsync(request, CancellationToken.None);
|
|
|
|
_ = mediator.Received(1).CreateStream(
|
|
Arg.Is<RAGStreamQuery>(query =>
|
|
query.KnowledgeBaseId == request.KnowledgeBaseId
|
|
&& query.Question == request.Question),
|
|
Arg.Any<CancellationToken>());
|
|
body.Position = 0;
|
|
var response = await new StreamReader(body).ReadToEndAsync();
|
|
var sourcesJson = JsonSerializer.Serialize(new { sources });
|
|
var contentJson = JsonSerializer.Serialize(new { content = "回答" });
|
|
response.Should().Be(
|
|
$"event: sources\ndata: {sourcesJson}\n\n" +
|
|
$"data: {contentJson}\n\n" +
|
|
"data: [DONE]\n\n");
|
|
endpoint.HttpContext.Response.ContentType.Should().Be("text/event-stream");
|
|
endpoint.HttpContext.Response.Headers.CacheControl.ToString().Should().Be("no-cache");
|
|
endpoint.HttpContext.Response.Headers.Connection.ToString().Should().Be("keep-alive");
|
|
}
|
|
|
|
private static async IAsyncEnumerable<RAGStreamEvent> CreateEvents(
|
|
IReadOnlyList<SourceChunk> sources)
|
|
{
|
|
await Task.Yield();
|
|
yield return RAGStreamEvent.FromSources(sources);
|
|
yield return RAGStreamEvent.Chunk("回答");
|
|
yield return RAGStreamEvent.Done();
|
|
}
|
|
}
|