AgenticCode/tests/ChatAgent.Api.Tests/ChatControllerTests.cs

using System.Net;
using System.Net.Http.Json;
using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json;
using ChatAgent.Shared.Models;
using Microsoft.AspNetCore.Mvc.Testing;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.ChatCompletion;
using Microsoft.SemanticKernel.Connectors.OpenAI;
using Moq;

namespace ChatAgent.Api.Tests;

public class ChatControllerTests : IClassFixture<WebApplicationFactory<Program>>
{
    private readonly WebApplicationFactory<Program> _factory;

    public ChatControllerTests(WebApplicationFactory<Program> factory)
    {
        _factory = factory;
    }

    [Fact]
    public async Task PostChat_StreamsTextDeltas_AndDone()
    {
        // Arrange: mock IChatCompletionService to return streaming text chunks
        var mockChatService = new Mock<IChatCompletionService>();

        var chunks = new List<StreamingChatMessageContent>
        {
            new(AuthorRole.Assistant, "Hello"),
            new(AuthorRole.Assistant, " world")
        };

        mockChatService
            .Setup(s => s.GetStreamingChatMessageContentsAsync(
                It.IsAny<ChatHistory>(),
                It.IsAny<PromptExecutionSettings>(),
                It.IsAny<Kernel>(),
                It.IsAny<CancellationToken>()))
            .Returns(chunks.ToAsyncEnumerable());

        var client = CreateClientWithMockedChatService(mockChatService.Object);

        var request = new ChatRequest
        {
            Messages = new List<ChatMessage>
            {
                new() { Role = "user", Content = "Hi" }
            }
        };

        // Act
        var response = await client.PostAsJsonAsync("/api/chat", request);

        Assert.Equal(HttpStatusCode.OK, response.StatusCode);
        Assert.Equal("text/event-stream", response.Content.Headers.ContentType?.MediaType);

        var body = await response.Content.ReadAsStringAsync();
        var events = ParseSSEData(body);

        // Assert: should contain text deltas and [DONE]
        Assert.Contains(events, e => e.Contains("\"text\":\"Hello\""));
        Assert.Contains(events, e => e.Contains("\"text\":\" world\""));
        Assert.Contains(events, e => e == "[DONE]");
    }

    [Fact]
    public async Task PostChat_HandlesServiceError_ReturnsErrorEvent()
    {
        // Arrange: mock IChatCompletionService to throw HttpRequestException
        var mockChatService = new Mock<IChatCompletionService>();

        mockChatService
            .Setup(s => s.GetStreamingChatMessageContentsAsync(
                It.IsAny<ChatHistory>(),
                It.IsAny<PromptExecutionSettings>(),
                It.IsAny<Kernel>(),
                It.IsAny<CancellationToken>()))
            .Returns(ThrowingAsyncEnumerable(new HttpRequestException("Connection refused")));

        var client = CreateClientWithMockedChatService(mockChatService.Object);

        var request = new ChatRequest
        {
            Messages = new List<ChatMessage>
            {
                new() { Role = "user", Content = "Hi" }
            }
        };

        // Act
        var response = await client.PostAsJsonAsync("/api/chat", request);
        var body = await response.Content.ReadAsStringAsync();
        var events = ParseSSEData(body);

        // Assert: should contain error and [DONE]
        Assert.Contains(events, e => e.Contains("\"error\""));
        Assert.Contains(events, e => e == "[DONE]");
    }

    private HttpClient CreateClientWithMockedChatService(IChatCompletionService chatService)
    {
        return _factory.WithWebHostBuilder(builder =>
        {
            builder.ConfigureServices(services =>
            {
                // Remove existing IChatCompletionService registration and replace with mock
                var descriptor = services.SingleOrDefault(
                    d => d.ServiceType == typeof(IChatCompletionService));
                if (descriptor != null)
                    services.Remove(descriptor);

                services.AddSingleton(chatService);
            });
        }).CreateClient();
    }

    /// <summary>
    /// Creates an IAsyncEnumerable that throws the given exception when iterated.
    /// Used to simulate service failures in streaming responses.
    /// </summary>
    private static async IAsyncEnumerable<StreamingChatMessageContent> ThrowingAsyncEnumerable(
        Exception exception,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await Task.CompletedTask;
        throw exception;
        yield break; // Required to make the compiler treat this as an async enumerable
    }

    [Fact]
    public async Task PostChat_ClarificationMessage_StreamedToClient()
    {
        // Arrange: simulate the LLM returning a clarification message
        // (what happens after the agent exhausts tool call retries)
        var mockChatService = new Mock<IChatCompletionService>();

        var clarificationText = "I couldn't determine the currency. Could you specify whether this is USD or GBP?";
        var chunks = new List<StreamingChatMessageContent>
        {
            new(AuthorRole.Assistant, clarificationText)
        };

        mockChatService
            .Setup(s => s.GetStreamingChatMessageContentsAsync(
                It.IsAny<ChatHistory>(),
                It.IsAny<PromptExecutionSettings>(),
                It.IsAny<Kernel>(),
                It.IsAny<CancellationToken>()))
            .Returns(chunks.ToAsyncEnumerable());

        var client = CreateClientWithMockedChatService(mockChatService.Object);

        var request = new ChatRequest
        {
            Messages = new List<ChatMessage>
            {
                new() { Role = "user", Content = "Invoice for 3 hours consulting" }
            }
        };

        // Act
        var response = await client.PostAsJsonAsync("/api/chat", request);
        var body = await response.Content.ReadAsStringAsync();
        var events = ParseSSEData(body);

        // Assert: clarification message streamed as text, not swallowed
        Assert.Contains(events, e => e.Contains("currency"));
        Assert.Contains(events, e => e == "[DONE]");
    }

    /// <summary>
    /// Extracts the data payloads from SSE-formatted text.
    /// </summary>
    private static List<string> ParseSSEData(string sseText)
    {
        return sseText.Split('\n')
            .Where(line => line.StartsWith("data: "))
            .Select(line => line.Substring(6).Trim())
            .ToList();
    }
}