"""测试辅助函数和工具类。 提供统一的测试辅助工具,减少测试代码重复。 """ import shutil from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable from urllib.parse import urlparse 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 get_bound_tcp_port(site: Any) -> int: """Resolve the bound aiohttp TCP site port for tests. We prefer the public ``site.name`` first. Some aiohttp test setups with ephemeral ports may not expose a usable port there, so we fall back to ``site._server.sockets`` as a test-only compatibility path. """ parsed = urlparse(getattr(site, "name", "")) if parsed.port is not None and parsed.port > 0: return parsed.port server = getattr(site, "_server", None) sockets = getattr(server, "sockets", None) if server else None if sockets: return sockets[0].getsockname()[1] raise RuntimeError("Unable to resolve bound TCP port from aiohttp site") # ============================================================ # 平台配置工厂 # ============================================================ 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