开发基于 Google Gemini API 的应用程序时,单元测试是保证代码质量的关键环节。然而,直接调用真实的 Gemini API 会引入网络延迟、费用、以及不确定的响应内容,使得测试不稳定且难以重复。为了解决这个问题,我们可以创建一个模拟客户端(MockClient),它完全兼容官方 google-genai SDK 的接口,使得业务代码无需任何修改即可在测试环境中运行。
本文介绍的 MockClient 为 Gemini API 的单元测试提供了一个强大而灵活的解决方案。它具有以下优点:
- 零侵入:业务代码与真实调用完全一致,无需任何修改。
- 全面覆盖:支持文本、流式、多轮对话、多模态图片和文件附件、思考等级配置以及函数调用等所有常见 API 用法。
- 高度可控:通过自定义响应生成器,可以精确控制每个测试的回复内容。
- 易于集成:可以与 pytest、unittest 等主流测试框架无缝配合。
通过使用 MockClient,开发者可以在隔离外部依赖的情况下,快速、稳定地测试与 Gemini API 交互的代码,从而提升开发效率和代码质量。
完整代码获取:上文中的 MockClient 实现即为完整代码,可直接复制到测试文件中使用。如果需要在生产代码中模拟更复杂的场景(例如模拟工具调用结果),可以在此基础上进一步扩展 MockModels._get_text 方法。
1. 为什么需要 MockClient?
在测试中,我们希望验证业务逻辑是否正确,而不是验证 API 是否正常工作。通过 MockClient,我们可以:
- 消除外部依赖:测试不依赖网络和 API 密钥,可以在任何环境中运行。
- 控制响应内容:可以返回固定文本或根据输入动态生成,便于验证不同分支。
- 提高测试速度:无需等待真实 API 响应,测试几乎瞬时完成。
- 保证一致性:测试结果可重复,不受模型更新或网络波动影响。
2. MockClient 的设计与实现
MockClient 模拟了 google.genai.Client 的核心接口,包括 models.generate_content、models.generate_content_stream 和 chats.create。同时,它能够解析输入中的各种数据类型(字符串、types.Part 对象、字节、文件路径等),并在模拟响应中体现图片或文件的存在。
以下是完整的 MockClient 实现代码:
from typing import Iterator, Union, List, Dict, Any, Optional, Callable, Tuple
import os
class MockResponse:
"""模拟 API 响应,仅包含 text 属性"""
def __init__(self, text: str):
self.text = text
class MockChat:
"""模拟聊天会话,支持多轮对话和多模态消息"""
def __init__(self, model: str, response_generator: Optional[Callable] = None):
self.model = model
self.history: List[str] = []
self._response_generator = response_generator
def send_message(self, message: Union[str, List[Any]]) –> MockResponse:
# 将消息序列化用于历史记录(仅取文本部分)
text_part = self._extract_text_from_message(message)
self.history.append(f"User: {text_part}")
# 生成回复
if self._response_generator:
has_image = self._check_image_in_message(message)
try:
reply = self._response_generator(
model=self.model,
contents=message,
is_chat=True,
has_image=has_image
)
except TypeError:
# 兼容旧版生成器(不接受 has_image 参数)
reply = self._response_generator(model=self.model, contents=message, is_chat=True)
else:
reply = self._default_generate(self.model, message, is_chat=True)
self.history.append(f"Assistant: {reply}")
return MockResponse(reply)
def _extract_text_from_message(self, message: Union[str, List[Any]]) –> str:
"""从消息中提取纯文本用于历史记录"""
if isinstance(message, str):
return message
parts = []
if isinstance(message, list):
for item in message:
if isinstance(item, str):
parts.append(item)
elif hasattr(item, 'text') and item.text:
parts.append(item.text)
elif isinstance(item, dict) and 'text' in item:
parts.append(item['text'])
return " ".join(parts) if parts else "[非文本内容]"
def _check_image_in_message(self, message: Union[str, List[Any]]) –> bool:
"""检查消息中是否包含图片或文件附件"""
if isinstance(message, list):
for item in message:
if hasattr(item, 'inline_data') and item.inline_data:
return True
if hasattr(item, 'file_data') and item.file_data:
return True
if isinstance(item, bytes):
return True
if isinstance(item, str) and os.path.isfile(item):
return True
return False
def _default_generate(self, model: str, contents: Any, is_chat: bool) –> str:
text, has_image = MockModels._extract_content_info(contents)
base = f"模拟聊天回复:关于“{text}”" if is_chat else f"模拟生成结果:针对“{text}”"
if has_image:
base += "(已接收文件内容)"
return base
class MockChats:
"""模拟 chats 模块"""
def __init__(self, response_generator: Optional[Callable] = None):
self._response_generator = response_generator
def create(self, model: str) –> MockChat:
return MockChat(model, self._response_generator)
class MockModels:
"""模拟 models 模块,支持多模态内容"""
def __init__(self, response_generator: Optional[Callable] = None, stream_chunk_size: int = 5):
self._response_generator = response_generator
self.stream_chunk_size = stream_chunk_size
@staticmethod
def _extract_content_info(contents: Union[str, List[Any], bytes, Any]) –> Tuple[str, bool]:
"""
从 contents 中提取文本信息和是否包含文件/图片。
返回 (text, has_file)
"""
has_file = False
text_parts = []
# 处理字符串
if isinstance(contents, str):
if os.path.isfile(contents):
has_file = True
text_parts.append(f"[文件: {os.path.basename(contents)}]")
else:
text_parts.append(contents)
# 处理字节数据(假设为图片或文件)
elif isinstance(contents, bytes):
has_file = True
text_parts.append("[二进制文件数据]")
# 处理列表
elif isinstance(contents, list):
for item in contents:
if isinstance(item, str):
if os.path.isfile(item):
has_file = True
text_parts.append(f"[文件: {os.path.basename(item)}]")
else:
text_parts.append(item)
elif isinstance(item, bytes):
has_file = True
text_parts.append("[二进制文件数据]")
elif isinstance(item, dict):
if 'text' in item and item['text']:
text_parts.append(item['text'])
if 'inline_data' in item or 'file_data' in item:
has_file = True
text_parts.append("[文件数据]")
else:
if hasattr(item, 'text') and item.text:
text_parts.append(item.text)
if hasattr(item, 'inline_data') and item.inline_data:
has_file = True
text_parts.append("[内联文件数据]")
if hasattr(item, 'file_data') and item.file_data:
has_file = True
text_parts.append("[文件附件数据]")
if hasattr(item, 'data') and hasattr(item, 'mime_type'):
has_file = True
text_parts.append(f"[{item.mime_type} 数据]")
else:
text_parts.append(str(contents))
text = " ".join(text_parts) if text_parts else ""
return text, has_file
def _get_text(self, model: str, contents: Any, is_chat: bool = False) –> str:
if self._response_generator:
text, has_file = self._extract_content_info(contents)
try:
return self._response_generator(
model=model,
contents=contents,
is_chat=is_chat,
has_file=has_file
)
except TypeError:
return self._response_generator(model=model, contents=contents, is_chat=is_chat)
else:
text, has_file = self._extract_content_info(contents)
base = f"模拟聊天回复:关于“{text}”" if is_chat else f"模拟生成结果:针对“{text}”"
if has_file:
base += "(已接收文件内容)"
return base
def generate_content(
self,
model: str,
contents: Union[str, List[Any], bytes, Any],
config: Optional[Any] = None
) –> MockResponse:
text = self._get_text(model, contents)
return MockResponse(text)
def generate_content_stream(
self,
model: str,
contents: Union[str, List[Any], bytes, Any],
config: Optional[Any] = None
) –> Iterator[MockResponse]:
full_text = self._get_text(model, contents)
for i in range(0, len(full_text), self.stream_chunk_size):
yield MockResponse(full_text[i:i+self.stream_chunk_size])
class MockClient:
"""
模拟 google.genai.Client,支持多模态输入。
用法与真实客户端完全一致,可通过 response_generator 自定义响应逻辑。
"""
def __init__(
self,
response_generator: Optional[Callable] = None,
stream_chunk_size: int = 5
):
"""
:param response_generator: 可选,自定义响应生成函数。
函数签名可以为:
def func(model: str, contents: Any, is_chat: bool = False, has_file: bool = False) -> str
如果不需要 has_file,也可以使用旧签名 def func(model: str, contents: Any, is_chat: bool = False) -> str
:param stream_chunk_size: 流式输出时每个片段的字符数。
"""
self.models = MockModels(response_generator, stream_chunk_size)
self.chats = MockChats(response_generator)
2.1 核心特性
- 接口完全兼容:MockClient 拥有 models 和 chats 属性,方法与官方 SDK 完全一致。
- 多模态识别:能够识别输入中的 types.Part 对象、字节数据、本地文件路径,并在模拟回复中标注“已接收文件内容”。
- 自定义响应生成器:通过 response_generator 参数传入一个回调函数,可以根据模型、输入内容和是否包含文件来动态生成响应,满足各种测试场景。
- 流式输出模拟:将完整响应按指定大小分片,模拟真实流式传输。
- 历史记录:MockChat 会记录对话历史,方便断言。
3. 在测试中替换真实客户端
使用 MockClient 的关键步骤是将业务代码中的 genai.Client 替换为 MockClient。推荐使用 pytest 的 monkeypatch 或 unittest.mock.patch 来实现,这样业务代码无需任何改动。
3.1 使用 monkeypatch(pytest)
import genai
from your_module import ask_gemini # 业务函数
def test_ask_gemini(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
result = ask_gemini("你好")
assert "模拟" in result
3.2 使用 patch(unittest)
import unittest
from unittest.mock import patch
import genai
from your_module import ask_gemini
class TestGemini(unittest.TestCase):
@patch('genai.Client', new=MockClient)
def test_ask_gemini(self):
result = ask_gemini("你好")
self.assertIn("模拟", result)
3.3 自定义响应
如果测试需要特定的回复内容,可以在替换时传入自定义的 response_generator:
def custom_response(model, contents, is_chat=False, has_file=False):
if "天气" in str(contents):
return "今天天气晴朗"
return "默认回复"
monkeypatch.setattr(genai, "Client", lambda: MockClient(response_generator=custom_response))
4. 各类调用语句示例
下面展示所有常见的调用类型,以及它们在 Mock 测试中的写法。请注意,业务代码与真实调用完全一致,只是在测试环境中替换了 Client 类。
4.1 基础文本生成(非流式)
业务代码:
from google import genai
def generate_text(prompt):
client = genai.Client()
response = client.models.generate_content(
model="gemini-3.6-flash",
contents=prompt
)
return response.text
测试代码:
def test_generate_text(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
result = generate_text("解释量子计算")
assert result.startswith("模拟生成结果")
4.2 流式输出
业务代码:
def stream_text(prompt):
client = genai.Client()
full_text = ""
for chunk in client.models.generate_content_stream(
model="gemini-3.5-flash",
contents=prompt
):
full_text += chunk.text
return full_text
测试代码:
def test_stream_text(monkeypatch):
monkeypatch.setattr(genai, "Client", lambda: MockClient(stream_chunk_size=4))
result = stream_text("写一首诗")
assert "模拟生成结果" in result
4.3 多轮对话
业务代码:
def chat_example():
client = genai.Client()
chat = client.chats.create(model="gemini-3.6-flash")
first = chat.send_message("你好!")
second = chat.send_message("我刚才说了什么?")
return first.text, second.text
测试代码:
def test_chat(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
first, second = chat_example()
assert "模拟聊天回复" in first
assert "模拟聊天回复" in second
4.4 配置思考等级
业务代码:
from google import genai
from google.genai import types
def generate_with_thinking(prompt):
client = genai.Client()
response = client.models.generate_content(
model="gemini-3.1-pro-preview",
contents=prompt,
config=types.GenerateContentConfig(
thinking_config=types.ThinkingConfig(
thinking_level=types.ThinkingLevel.HIGH
)
)
)
return response.text
测试代码:
def test_with_thinking(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
result = generate_with_thinking("复杂问题")
assert "模拟生成结果" in result
4.5 多模态:发送图片(字节数据)
业务代码:
def analyze_image(image_bytes):
client = genai.Client()
response = client.models.generate_content(
model="gemini-3.5-flash",
contents=[
types.Part.from_text("描述这张图片"),
types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")
]
)
return response.text
测试代码:
def test_analyze_image(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
fake_image = b'\\xff\\xd8\\xff\\xe0fake'
result = analyze_image(fake_image)
assert "描述这张图片" in result
assert "文件" in result # 模拟回复会标注已接收文件
4.6 多模态:发送图片(文件路径)
业务代码:
def describe_image_file(path):
client = genai.Client()
response = client.models.generate_content(
model="gemini-3.6-flash",
contents=[
"这张图片里有什么?",
path
]
)
return response.text
测试代码:
def test_describe_image_file(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
result = describe_image_file("test_image.png")
assert "文件" in result
4.7 文件附件(例如 PDF)
业务代码:
def summarize_pdf(pdf_bytes):
client = genai.Client()
response = client.models.generate_content(
model="gemini-3.6-flash",
contents=[
types.Part.from_text("总结这份 PDF"),
types.Part.from_bytes(data=pdf_bytes, mime_type="application/pdf")
]
)
return response.text
测试代码:
def test_summarize_pdf(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
fake_pdf = b'%PDF-1.4 fake'
result = summarize_pdf(fake_pdf)
assert "文件" in result
4.8 文件附件(通过 URI)
业务代码:
def analyze_file_uri(file_uri, mime_type):
client = genai.Client()
response = client.models.generate_content(
model="gemini-3.6-flash",
contents=[
types.Part.from_uri(file_uri=file_uri, mime_type=mime_type),
"总结这个文件"
]
)
return response.text
测试代码:
def test_analyze_file_uri(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
uri = "https://generativelanguage.googleapis.com/v1beta/files/abc123"
result = analyze_file_uri(uri, "application/pdf")
assert "文件" in result
4.9 函数调用(Function Calling)
业务代码:
def get_current_weather(location: str) –> str:
return f"{location}的天气是晴天"
def ask_weather():
client = genai.Client()
response = client.models.generate_content(
model='gemini-3.5-flash',
contents='波士顿的天气怎么样?',
config=types.GenerateContentConfig(tools=[get_current_weather])
)
return response.text
测试代码:
def test_ask_weather(monkeypatch):
monkeypatch.setattr(genai, "Client", MockClient)
result = ask_weather()
assert "模拟生成结果" in result
5. 自定义响应生成器的高级用法
MockClient 允许通过 response_generator 实现非常灵活的响应控制。例如,我们可以根据输入内容返回不同的模拟结果:
def smart_generator(model, contents, is_chat=False, has_file=False):
# 将 contents 转为字符串便于判断
content_str = str(contents)
if "天气" in content_str:
return "今天天气晴朗,适合出行。"
if "图片" in content_str or has_file:
return "这是一张包含丰富内容的图片/文件。"
if is_chat:
return "这是对您消息的友好回复。"
return "这是通用模拟回复。"
monkeypatch.setattr(genai, "Client", lambda: MockClient(response_generator=smart_generator))
这样,不同的测试用例可以获得针对性的回复,便于验证业务逻辑的分支。






