File size: 16,726 Bytes
8ede856 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 | """测试辅助函数和工具类。
提供统一的测试辅助工具,减少测试代码重复。
"""
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable
from unittest.mock import AsyncMock, MagicMock
from astrbot.core.message.components import BaseMessageComponent
class NoopAwaitable:
"""可等待的空操作对象。
用于 mock 需要返回 awaitable 对象的方法。
"""
def __await__(self):
if False:
yield
return None
# ============================================================
# 平台配置工厂
# ============================================================
def make_platform_config(platform_type: str, **kwargs) -> dict:
"""平台配置工厂函数。
Args:
platform_type: 平台类型 (telegram, discord, aiocqhttp 等)
**kwargs: 覆盖默认配置的字段
Returns:
dict: 平台配置字典
"""
configs = {
"telegram": {
"id": "test_telegram",
"telegram_token": "test_token_123",
"telegram_api_base_url": "https://api.telegram.org/bot",
"telegram_file_base_url": "https://api.telegram.org/file/bot",
"telegram_command_register": True,
"telegram_command_auto_refresh": True,
"telegram_command_register_interval": 300,
"telegram_media_group_timeout": 2.5,
"telegram_media_group_max_wait": 10.0,
"start_message": "Welcome to AstrBot!",
},
"discord": {
"id": "test_discord",
"discord_token": "test_token_123",
"discord_proxy": None,
"discord_command_register": True,
"discord_guild_id_for_debug": None,
"discord_activity_name": "Playing AstrBot",
},
"aiocqhttp": {
"id": "test_aiocqhttp",
"ws_reverse_host": "0.0.0.0",
"ws_reverse_port": 6199,
"ws_reverse_token": "test_token",
},
"webchat": {
"id": "test_webchat",
},
"wecom": {
"id": "test_wecom",
"wecom_corpid": "test_corpid",
"wecom_secret": "test_secret",
},
}
config = configs.get(platform_type, {"id": f"test_{platform_type}"}).copy()
config.update(kwargs)
return config
# ============================================================
# Telegram 辅助函数
# ============================================================
def create_mock_update(
message_text: str | None = "Hello World",
chat_type: str = "private",
chat_id: int = 123456789,
user_id: int = 987654321,
username: str = "test_user",
message_id: int = 1,
media_group_id: str | None = None,
photo: list | None = None,
video: MagicMock | None = None,
document: MagicMock | None = None,
voice: MagicMock | None = None,
sticker: MagicMock | None = None,
reply_to_message: MagicMock | None = None,
caption: str | None = None,
entities: list | None = None,
caption_entities: list | None = None,
message_thread_id: int | None = None,
is_topic_message: bool = False,
):
"""创建模拟的 Telegram Update 对象。
Args:
message_text: 消息文本
chat_type: 聊天类型
chat_id: 聊天 ID
user_id: 用户 ID
username: 用户名
message_id: 消息 ID
media_group_id: 媒体组 ID
photo: 图片列表
video: 视频对象
document: 文档对象
voice: 语音对象
sticker: 贴纸对象
reply_to_message: 回复的消息
caption: 说明文字
entities: 实体列表
caption_entities: 说明实体列表
message_thread_id: 消息线程 ID
is_topic_message: 是否为主题消息
Returns:
MagicMock: 模拟的 Update 对象
"""
update = MagicMock()
update.update_id = 1
# Create message mock
message = MagicMock()
message.message_id = message_id
message.chat = MagicMock()
message.chat.id = chat_id
message.chat.type = chat_type
message.message_thread_id = message_thread_id
message.is_topic_message = is_topic_message
# Create user mock
from_user = MagicMock()
from_user.id = user_id
from_user.username = username
message.from_user = from_user
# Set message content
message.text = message_text
message.media_group_id = media_group_id
message.photo = photo
message.video = video
message.document = document
message.voice = voice
message.sticker = sticker
message.reply_to_message = reply_to_message
message.caption = caption
message.entities = entities
message.caption_entities = caption_entities
update.message = message
update.effective_chat = message.chat
return update
def create_mock_file(file_path: str = "https://api.telegram.org/file/test.jpg"):
"""创建模拟的 Telegram File 对象。
Args:
file_path: 文件路径
Returns:
MagicMock: 模拟的 File 对象
"""
file = MagicMock()
file.file_path = file_path
file.get_file = AsyncMock(return_value=file)
return file
# ============================================================
# Discord 辅助函数
# ============================================================
def create_mock_discord_attachment(
filename: str = "test.txt",
url: str = "https://cdn.discordapp.com/test.txt",
content_type: str | None = None,
size: int = 1024,
):
"""创建模拟的 Discord Attachment 对象。
Args:
filename: 文件名
url: 文件 URL
content_type: 内容类型
size: 文件大小
Returns:
MagicMock: 模拟的 Attachment 对象
"""
attachment = MagicMock()
attachment.filename = filename
attachment.url = url
attachment.content_type = content_type
attachment.size = size
return attachment
def create_mock_discord_user(
user_id: int = 123456789,
name: str = "TestUser",
display_name: str = "Test User",
bot: bool = False,
):
"""创建模拟的 Discord User 对象。
Args:
user_id: 用户 ID
name: 用户名
display_name: 显示名
bot: 是否为机器人
Returns:
MagicMock: 模拟的 User 对象
"""
user = MagicMock()
user.id = user_id
user.name = name
user.display_name = display_name
user.bot = bot
user.mention = f"<@{user_id}>"
return user
def create_mock_discord_channel(
channel_id: int = 111222333,
channel_type: str = "text",
name: str = "general",
guild_id: int | None = 444555666,
):
"""创建模拟的 Discord Channel 对象。
Args:
channel_id: 频道 ID
channel_type: 频道类型
name: 频道名
guild_id: 服务器 ID
Returns:
MagicMock: 模拟的 Channel 对象
"""
channel = MagicMock()
channel.id = channel_id
channel.name = name
channel.type = channel_type
if guild_id:
channel.guild = MagicMock()
channel.guild.id = guild_id
else:
channel.guild = None
return channel
# ============================================================
# 消息组件辅助函数
# ============================================================
def create_mock_message_component(
component_type: str,
**kwargs: Any,
) -> BaseMessageComponent:
"""创建模拟的消息组件。
Args:
component_type: 组件类型 (plain, image, at, reply, file)
**kwargs: 组件参数
Returns:
BaseMessageComponent: 消息组件实例
"""
from astrbot.core.message import components as Comp
component_map = {
"plain": Comp.Plain,
"image": Comp.Image,
"at": Comp.At,
"reply": Comp.Reply,
"file": Comp.File,
}
component_class = component_map.get(component_type.lower())
if not component_class:
raise ValueError(f"Unknown component type: {component_type}")
return component_class(**kwargs)
def create_mock_llm_response(
completion_text: str = "Hello! How can I help you?",
role: str = "assistant",
tools_call_name: list[str] | None = None,
tools_call_args: list[dict] | None = None,
tools_call_ids: list[str] | None = None,
):
"""创建模拟的 LLM 响应。
Args:
completion_text: 完成文本
role: 角色
tools_call_name: 工具调用名称列表
tools_call_args: 工具调用参数列表
tools_call_ids: 工具调用 ID 列表
Returns:
LLMResponse: 模拟的 LLM 响应
"""
from astrbot.core.provider.entities import LLMResponse, TokenUsage
return LLMResponse(
role=role,
completion_text=completion_text,
tools_call_name=tools_call_name or [],
tools_call_args=tools_call_args or [],
tools_call_ids=tools_call_ids or [],
usage=TokenUsage(input_other=10, output=5),
)
# ============================================================
# 测试插件辅助函数
# ============================================================
@dataclass
class MockPluginConfig:
"""测试插件配置。
用于创建和管理测试用的模拟插件。
Attributes:
name: 插件名称
author: 作者
description: 描述
version: 版本
repo: 仓库 URL
main_code: main.py 的代码内容
requirements: 依赖列表
has_readme: 是否创建 README.md
readme_content: README.md 内容
"""
name: str = "test_plugin"
author: str = "Test Author"
description: str = "A test plugin for unit testing"
version: str = "1.0.0"
repo: str = "https://github.com/test/test_plugin"
main_code: str = ""
requirements: list[str] = field(default_factory=list)
has_readme: bool = True
readme_content: str = "# Test Plugin\n\nThis is a test plugin."
# 默认的插件主代码模板
DEFAULT_PLUGIN_MAIN_TEMPLATE = '''
from astrbot.api import star
class Main(star.Star):
"""测试插件主类。"""
def __init__(self, context):
super().__init__(context)
self.name = "{plugin_name}"
async def initialize(self):
"""初始化插件。"""
pass
async def terminate(self):
"""终止插件。"""
pass
'''
class MockPluginBuilder:
"""测试插件构建器。
用于创建、管理和清理测试用的模拟插件。支持任意插件的模拟创建。
Example:
# 创建一个简单的测试插件
builder = MockPluginBuilder(plugin_store_path)
plugin_dir = builder.create("my_test_plugin")
# 创建自定义配置的插件
config = MockPluginConfig(
name="custom_plugin",
version="2.0.0",
main_code="print('hello')",
)
plugin_dir = builder.create(config)
# 清理插件
builder.cleanup("my_test_plugin")
"""
def __init__(self, plugin_store_path: str | Path):
"""初始化构建器。
Args:
plugin_store_path: 插件存储路径 (通常是 data/plugins)
"""
self.plugin_store_path = Path(plugin_store_path)
self._created_plugins: set[str] = set()
def create(
self,
plugin_config: str | MockPluginConfig | None = None,
**kwargs,
) -> Path:
"""创建模拟插件。
Args:
plugin_config: 插件名称字符串、MockPluginConfig 对象或 None
**kwargs: 如果 plugin_config 是字符串或 None,这些参数用于构建 MockPluginConfig
Returns:
Path: 创建的插件目录路径
"""
# 处理不同类型的输入
if plugin_config is None:
config = MockPluginConfig(**kwargs)
elif isinstance(plugin_config, str):
config = MockPluginConfig(name=plugin_config, **kwargs)
elif isinstance(plugin_config, MockPluginConfig):
config = plugin_config
else:
raise TypeError(f"Invalid plugin_config type: {type(plugin_config)}")
# 创建插件目录
plugin_dir = self.plugin_store_path / config.name
plugin_dir.mkdir(parents=True, exist_ok=True)
# 创建 metadata.yaml
metadata_content = "\n".join(
[
f"name: {config.name}",
f"author: {config.author}",
f"desc: {config.description}",
f"version: {config.version}",
f"repo: {config.repo}",
]
)
(plugin_dir / "metadata.yaml").write_text(
metadata_content + "\n", encoding="utf-8"
)
# 创建 main.py
main_code = config.main_code or DEFAULT_PLUGIN_MAIN_TEMPLATE.format(
plugin_name=config.name
)
(plugin_dir / "main.py").write_text(main_code, encoding="utf-8")
# 创建 requirements.txt(如果有依赖)
if config.requirements:
(plugin_dir / "requirements.txt").write_text(
"\n".join(config.requirements) + "\n", encoding="utf-8"
)
# 创建 README.md(如果需要)
if config.has_readme:
(plugin_dir / "README.md").write_text(
config.readme_content, encoding="utf-8"
)
# 记录创建的插件
self._created_plugins.add(config.name)
return plugin_dir
def cleanup(self, plugin_name: str | None = None) -> None:
"""清理插件。
Args:
plugin_name: 要清理的插件名称,如果为 None 则清理所有由本构建器创建的插件
"""
if plugin_name:
plugins_to_clean = {plugin_name}
else:
plugins_to_clean = self._created_plugins.copy()
for name in plugins_to_clean:
plugin_dir = self.plugin_store_path / name
if plugin_dir.exists():
shutil.rmtree(plugin_dir)
self._created_plugins.discard(name)
def cleanup_all(self) -> None:
"""清理所有由本构建器创建的插件。"""
self.cleanup(None)
def get_plugin_path(self, plugin_name: str) -> Path:
"""获取插件路径。
Args:
plugin_name: 插件名称
Returns:
Path: 插件目录路径
"""
return self.plugin_store_path / plugin_name
@property
def created_plugins(self) -> set[str]:
"""获取已创建的插件名称集合。"""
return self._created_plugins.copy()
def create_mock_updater_install(
plugin_builder: MockPluginBuilder,
repo_to_plugin: dict[str, str] | None = None,
) -> Callable:
"""创建模拟的 updater.install 方法。
Args:
plugin_builder: MockPluginBuilder 实例
repo_to_plugin: 仓库 URL 到插件名称的映射,格式: {"https://github.com/user/repo": "plugin_name"}
Returns:
Callable: 异步函数,可用于 monkeypatch.setattr
"""
async def mock_install(repo_url: str, proxy: str = "") -> str:
"""Mock updater.install 方法。"""
# 查找插件名称
plugin_name = None
if repo_to_plugin:
plugin_name = repo_to_plugin.get(repo_url)
# 如果没有映射,尝试从 URL 提取插件名
if not plugin_name:
# 从 https://github.com/user/plugin_name 提取 plugin_name
parts = repo_url.rstrip("/").split("/")
plugin_name = parts[-1] if parts else "unknown_plugin"
# 创建插件目录
config = MockPluginConfig(name=plugin_name, repo=repo_url)
plugin_dir = plugin_builder.create(config)
return str(plugin_dir)
return mock_install
def create_mock_updater_update(
plugin_builder: MockPluginBuilder,
update_callback: Callable | None = None,
) -> Callable:
"""创建模拟的 updater.update 方法。
Args:
plugin_builder: MockPluginBuilder 实例
update_callback: 更新回调函数,接收 plugin 参数
Returns:
Callable: 异步函数,可用于 monkeypatch.setattr
"""
async def mock_update(plugin, proxy: str = "") -> None:
"""Mock updater.update 方法。"""
plugin_dir = plugin_builder.get_plugin_path(plugin.name)
# 创建更新标记文件
(plugin_dir / ".updated").write_text("ok", encoding="utf-8")
# 调用回调
if update_callback:
update_callback(plugin)
return mock_update
|