llm-pii-poc/tests/PiiRedaction.Infrastructure.Tests/Llm/MockLlmPromptServiceTests.cs

using System.Runtime.CompilerServices;
using FluentAssertions;
using Microsoft.Extensions.AI;
using PiiRedaction.Infrastructure.Llm;

namespace PiiRedaction.Infrastructure.Tests.Llm;

[TestFixture]
public sealed class MockLlmPromptServiceTests
{
    [Test]
    public async Task SendPromptAsync_ReturnsAssistantText()
    {
        var service = new MockLlmPromptService(new MockChatClient());

        const string prompt = "Sanitized <EMAIL_1> content.";
        var response = await service.SendPromptAsync(prompt);

        response.Should().Contain("Mock LLM Response");
        response.Should().Contain($"{prompt.Length} chars");
    }

    [Test]
    public async Task SendPromptAsync_UsesUserRoleForOutboundMessage()
    {
        var capturingClient = new CapturingChatClient();
        var service = new MockLlmPromptService(capturingClient);

        await service.SendPromptAsync("Only placeholders <PHONE_1>.");

        capturingClient.LastUserMessage.Should().Be("Only placeholders <PHONE_1>.");
    }

    private sealed class CapturingChatClient : IChatClient
    {
        public string? LastUserMessage { get; private set; }

        public Task<ChatResponse> GetResponseAsync(
            IEnumerable<ChatMessage> messages,
            ChatOptions? options = null,
            CancellationToken cancellationToken = default)
        {
            LastUserMessage = messages.LastOrDefault(message => message.Role == ChatRole.User)?.Text;
            return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")));
        }

        public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
            IEnumerable<ChatMessage> messages,
            ChatOptions? options = null,
            [EnumeratorCancellation] CancellationToken cancellationToken = default)
        {
            await GetResponseAsync(messages, options, cancellationToken);
            yield break;
        }

        public object? GetService(Type serviceType, object? serviceKey = null) => null;

        public void Dispose()
        {
        }
    }
}