Azure.AI.OpenAI のクラスのモックを作りたい

Azure OpenAI Service を .NET から呼び出す時には Azure.AI.OpenAI パッケージを使うのが一般的です。
このパッケージを使っているときに ChatClient などのクライアントを単体テスト用にモックにするときに少し困ったことがあったのでメモしておきます。

モックの作り方

モックを作ること自体は簡単です。ChatClient の各メソッドは、以下のように virtual がついていてモックが作りやすいようになっています。

    public virtual async Task<ClientResult<ChatCompletion>> CompleteChatAsync(IEnumerable<ChatMessage> messages, ChatCompletionOptions options = null, CancellationToken cancellationToken = default(CancellationToken))

このようなメソッドを持つクラスをモックするには、Moq などのモックライブラリを使って以下のように書けば良いとおもっていました。

// このコードは動きません
var mock = new Mock<ChatClient>();
mock.Setup(x => x.CompleteChatAsync(It.IsAny<IEnumerable<ChatMessage>>(), It.IsAny<ChatCompletionOptions>(), It.IsAny<CancellationToken>()))
    .ReturnsAsync(new ClientResult<ChatCompletion>(new ChatCompletion()));

問題点と解決策

各メソッドは virtual で、引数のオブジェクトは普通にインスタンス化出来るのですが戻り値の型の ChatCompletion のコンストラクタが軒並み internal になっているためモック用の戻り値を作ることができません。どのように作るのかわからなかったのですが専用のファクトリクラスが用意されていました。

OpenAI.Chat.OpenAIChatModelFactory というクラスでチャットに関するモデルクラスを作成するためのメソッドが用意されています。
例えば CompleteChatAsync の戻り値で Mock!! というテキストを返す場合は以下のように書けます。

using Moq;
using OpenAI.Chat;
using OpenAI.Models;
using System.ClientModel;
using System.ClientModel.Primitives;
using System.Threading.Tasks;

namespace TestProject2;

[TestClass]
public sealed class Test1
{
    [TestMethod]
    public async Task TestMethod1()
    {
        var mockClient = new Mock<ChatClient>();
        mockClient.Setup(x => x.CompleteChatAsync(It.IsAny<ChatMessage[]>()))
            .ReturnsAsync(() => 
                // ClientResult は FromValue メソッドで作る
                ClientResult.FromValue(
                    // ChatCompletion は OpenAIChatModelFactory.ChatCompletion で作る
                    OpenAIChatModelFactory.ChatCompletion(content: new ("Mock!!")), 
                    // PipelineResponse はテストのアサートでは特に使用しないので純粋なモックにする
                    Mock.Of<PipelineResponse>()));

        var response = await mockClient.Object.CompleteChatAsync([]);
        Assert.AreEqual("Mock!!", response.Value.Content[0].Text);
    }
}

引数もアサートする場合は通常の Moq の使い方と同じように It クラスを使って以下のように書けます。

using Moq;
using OpenAI.Chat;
using System.ClientModel;
using System.ClientModel.Primitives;

namespace TestProject2;

[TestClass]
public sealed class Test1
{
    [TestMethod]
    public async Task TestMethod1()
    {
        var mockClient = new Mock<ChatClient>();
        mockClient.Setup(x => x.CompleteChatAsync(It.IsAny<ChatMessage[]>()))
            .ReturnsAsync(() =>
                ClientResult.FromValue(
                    OpenAIChatModelFactory.ChatCompletion(content: new("Mock!!")),
                    Mock.Of<PipelineResponse>()));
        
        var response = await mockClient.Object.CompleteChatAsync([new UserChatMessage("Hello")]);
        Assert.AreEqual("Mock!!", response.Value.Content[0].Text);
        // 期待する引数で呼び出されたかどうかを確認
        mockClient.Verify(x => x.CompleteChatAsync(
            It.Is<ChatMessage[]>(x => x.Length == 1 && x[0].Content[0].Text == "Hello")), 
            Times.Once);
    }
}

ToolCall を伴うようなものの場合は以下のように書けます。

using Moq;
using OpenAI.Chat;
using System.ClientModel;
using System.ClientModel.Primitives;

namespace TestProject2;

[TestClass]
public sealed class Test1
{
    [TestMethod]
    public async Task TestMethod1()
    {
        var mockClient = new Mock<ChatClient>();
        mockClient.Setup(x => x.CompleteChatAsync(It.IsAny<ChatMessage[]>()))
            .ReturnsAsync(() =>
                ClientResult.FromValue(
                    OpenAIChatModelFactory.ChatCompletion(
                        finishReason: ChatFinishReason.ToolCalls,
                        toolCalls: [ 
                            ChatToolCall.CreateFunctionToolCall(
                                "id1", 
                                "func", 
                                BinaryData.FromObjectAsJson(new { name = "tanaka" }) ) 
                        ]),
                    Mock.Of<PipelineResponse>()));
        
        var response = await mockClient.Object.CompleteChatAsync([new UserChatMessage("Hello")]);
        // ツール呼び出し
        Assert.AreEqual(ChatFinishReason.ToolCalls, response.Value.FinishReason);
        // ツール呼び出しの内容
        Assert.AreEqual(1, response.Value.ToolCalls.Count);
        var tool = response.Value.ToolCalls[0];
        Assert.AreEqual("id1", tool.Id);
        Assert.AreEqual("func", tool.FunctionName);
        Assert.AreEqual("{\"name\":\"tanaka\"}", tool.FunctionArguments.ToString());
        // 期待する引数で呼び出されたかどうかを確認
        mockClient.Verify(x => x.CompleteChatAsync(
            It.Is<ChatMessage[]>(x => x.Length == 1 && x[0].Content[0].Text == "Hello")), 
            Times.Once);
    }
}

まとめ

Azure.AI.OpenAI (本家の OpenAI も同じです) の ChatCompletion クラスは new でインスタンス化できませんが、ファクトリーメソッドを使うことでモックを作ることができました。
実際には、クライアントクラスを直接モックするよりも、1枚レイヤーかませて XxxxxRepository などのクラスを作ってそちらをモックする方が良いかもしれませんが、一応 ChatClient などのクラスをモックする方法もあることを覚えておくと良いかもしれません。