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
などのクラスをモックする方法もあることを覚えておくと良いかもしれません。
Discussion