Mocking CompleteChatStreamingAsync

[ethan-christensen-cag]

Hello all,

I am just reaching out to see if there is any experience with mocking CompleteChatStreamingAsync with OpenAI Chat. I am wanting to test a service that utilizes OpenAI ChatClients CompleteChatStreamingAsync, though I can never seem to make it work.

My service consumes an instance of ChatClient which I am trying to Mock.

[Fact]
    public async Task StreamChatAsync_ShouldReturnStreamedContent()
    {
        // Arrange
        List<StreamingChatCompletionUpdate> mockStreamingChatCompletionUpdates =
        [
            OpenAIChatModelFactory.StreamingChatCompletionUpdate(contentUpdate: ["A"]),
            OpenAIChatModelFactory.StreamingChatCompletionUpdate(contentUpdate: ["B"]),
        ];
        
        _mockChatClient.Setup(client => client.CompleteChatStreamingAsync(It.IsAny<ChatMessage[]>(),
            It.IsAny<ChatCompletionOptions>(), 
            It.IsAny<CancellationToken>())
        ).Returns(mockStreamingChatCompletionUpdates.ToAsyncEnumerable());

        // Act
        var totalChunks = 0;
        var accumulatedContent = "";
        
        await foreach (var chunk in _openAiService.StreamChatAsync("Hello"))
        {
            totalChunks++;
            accumulatedContent += chunk;
        }
        
        // Assert
        totalChunks.Should().Be(2);
        accumulatedContent.Should().Be("AB");
    }

But am always having trouble mocking the streaming response with error:

Cannot resolve method:
    Returns(IAsyncEnumerable<StreamingChatCompletionUpdate>)
Candidates are:
    Returns(AsyncCollectionResult<StreamingChatCompletionUpdate>)
    Returns(Delegate)
    Returns(Func<AsyncCollectionResult<StreamingChatCompletionUpdate>>)
    Returns(InvocationFunc)
    all from interface IReturns<TMock,TResult>

with this section:

_mockChatClient.Setup(client => client.CompleteChatStreamingAsync(It.IsAny<ChatMessage[]>(),
            It.IsAny<ChatCompletionOptions>(), 
            It.IsAny<CancellationToken>())
        ).Returns(mockStreamingChatCompletionUpdates.ToAsyncEnumerable());

[ethan-christensen-cag]

For the sake of anyone else trying to Mock such a thing, I was able to solve my issue by creating my own class which implements AsyncCollectionResult. Then using my implementation I could essentially create a AsyncCollectionResult object which I could just pass to the setup for Open AI Chat Client.

Implemntation

public class MockAsyncCollectionResult<T>(IEnumerable<T> items) : AsyncCollectionResult<T>
{
    public override async IAsyncEnumerable<ClientResult<T>> GetRawPagesAsync()
    {
        foreach (var item in items)
        {
            var fakePipelineResponse = new Mock<PipelineResponse>().Object;
            yield return ClientResult.FromValue(item, fakePipelineResponse);
        }
        await Task.CompletedTask;
    }

    protected override async IAsyncEnumerable<T> GetValuesFromPageAsync(ClientResult page)
    {
        var clientResult = (ClientResult<T>)page;
        yield return clientResult.Value;
        await Task.CompletedTask;
    }

    public override ContinuationToken? GetContinuationToken(ClientResult page)
    {
        return null;
    }
}

Sample Usage

[Fact]
public async Task StreamChatAsync_ShouldReturnStreamedContent()
{
    // Arrange
    var streamingUpdates= new List<StreamingChatCompletionUpdate>
    {
        OpenAIChatModelFactory.StreamingChatCompletionUpdate(contentUpdate: ["The color of the"]),
        OpenAIChatModelFactory.StreamingChatCompletionUpdate(contentUpdate: [" sky is generally blue. "]),
        OpenAIChatModelFactory.StreamingChatCompletionUpdate(contentUpdate: ["Let me know if you "]),
        OpenAIChatModelFactory.StreamingChatCompletionUpdate(contentUpdate: [" have any further questions."]),
    };
    
    _mockChatClient.Setup(x => x.CompleteChatStreamingAsync(It.IsAny<IEnumerable<ChatMessage>>(),
            It.IsAny<ChatCompletionOptions>(), It.IsAny<CancellationToken>()))
        .Returns(new MockAsyncCollectionResult<StreamingChatCompletionUpdate>(streamingUpdates));

    // Act & Assert
    var totalChunks = 0;

    await foreach (var chunk in _serviceYouAreTesting.StreamChatAsync("Hello"))
    {
        totalChunks++;
        // Do what ever else you need to validate your service returned the right information.
    }

    totalChunks.Should().Be(4);
}