diff --git a/tests/unit/test_agent_message.py b/tests/unit/test_agent_message.py new file mode 100644 index 000000000..b205a23f5 --- /dev/null +++ b/tests/unit/test_agent_message.py @@ -0,0 +1,280 @@ +"""Unit tests for astrbot.core.agent.message: Message, CheckpointData, CheckpointMessageSegment.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from astrbot.core.agent.message import ( + AssistantMessageSegment, + CheckpointData, + CheckpointMessageSegment, + ContentPart, + ImageURLPart, + Message, + SystemMessageSegment, + TextPart, + ThinkPart, + ToolCall, + ToolCallMessageSegment, + UserMessageSegment, + bind_checkpoint_messages, + dump_messages_with_checkpoints, + get_checkpoint_id, + is_checkpoint_message, + strip_checkpoint_messages, +) + + +class TestMessageConstruction: + """Message base model construction and validation.""" + + def test_minimal_user_message(self): + """A user message requires role and content.""" + msg = Message(role="user", content="hello") + assert msg.role == "user" + assert msg.content == "hello" + assert msg.tool_calls is None + assert msg.tool_call_id is None + + def test_assistant_message_with_content(self): + """An assistant message with text content.""" + msg = Message(role="assistant", content="I am an assistant") + assert msg.role == "assistant" + assert msg.content == "I am an assistant" + + def test_assistant_message_with_tool_calls_no_content(self): + """Assistant messages with tool_calls may have content=None.""" + tc = ToolCall(id="call_1", function=ToolCall.FunctionBody(name="f", arguments="{}")) + msg = Message(role="assistant", content=None, tool_calls=[tc]) + assert msg.role == "assistant" + assert msg.content is None + assert len(msg.tool_calls or []) == 1 + + def test_tool_message(self): + """A tool message with role='tool'.""" + msg = Message(role="tool", content="tool result", tool_call_id="call_1") + assert msg.role == "tool" + assert msg.tool_call_id == "call_1" + + def test_system_message(self): + """A system message.""" + msg = Message(role="system", content="You are a helpful assistant") + assert msg.role == "system" + + def test_missing_content_raises_for_user(self): + """User/System/Tool messages must have content.""" + with pytest.raises(ValidationError, match="content is required"): + Message(role="user", content=None) + + def test_missing_content_raises_for_system(self): + """System messages must have content.""" + with pytest.raises(ValidationError, match="content is required"): + Message(role="system", content=None) + + def test_invalid_role_raises(self): + """An invalid role is rejected.""" + with pytest.raises(ValidationError): + Message(role="invalid_role", content="hi") + + +class TestCheckpointMessage: + """CheckpointData and role='_checkpoint' messages.""" + + def test_checkpoint_data_construction(self): + """CheckpointData can be constructed with an id.""" + cp = CheckpointData(id="cp_1") + assert cp.id == "cp_1" + + def test_checkpoint_message_valid(self): + """A valid checkpoint message has role _checkpoint and CheckpointData content.""" + cp = CheckpointData(id="cp_1") + msg = Message(role="_checkpoint", content=cp) + assert msg.role == "_checkpoint" + assert isinstance(msg.content, CheckpointData) + + def test_checkpoint_message_string_content_raises(self): + """Checkpoint messages must use CheckpointData, not plain strings.""" + with pytest.raises(ValidationError): + Message(role="_checkpoint", content="not a checkpoint") + + def test_checkpoint_data_in_non_checkpoint_role_raises(self): + """CheckpointData in a non-checkpoint role is rejected.""" + cp = CheckpointData(id="cp_1") + with pytest.raises(ValidationError, match="CheckpointData is only allowed"): + Message(role="user", content=cp) + + def test_is_checkpoint_message_detection(self): + """is_checkpoint_message correctly identifies checkpoint messages.""" + cp = CheckpointData(id="cp_1") + msg = Message(role="_checkpoint", content=cp) + assert is_checkpoint_message(msg) is True + assert is_checkpoint_message(Message(role="user", content="hi")) is False + + def test_is_checkpoint_message_dict(self): + """is_checkpoint_message works with dicts.""" + assert is_checkpoint_message({"role": "_checkpoint"}) is True + assert is_checkpoint_message({"role": "user"}) is False + + def test_get_checkpoint_id(self): + """get_checkpoint_id returns the id from a checkpoint message.""" + cp = CheckpointData(id="cp_42") + msg = Message(role="_checkpoint", content=cp) + assert get_checkpoint_id(msg) == "cp_42" + + def test_get_checkpoint_id_none_for_non_checkpoint(self): + """get_checkpoint_id returns None for non-checkpoint messages.""" + msg = Message(role="user", content="hi") + assert get_checkpoint_id(msg) is None + + def test_strip_checkpoint_messages(self): + """strip_checkpoint_messages removes checkpoint entries.""" + history = [ + {"role": "user", "content": "hi"}, + {"role": "_checkpoint", "content": {"id": "cp_1"}}, + {"role": "assistant", "content": "hello"}, + ] + cleaned = strip_checkpoint_messages(history) + assert len(cleaned) == 2 + assert all(m["role"] != "_checkpoint" for m in cleaned) + + def test_bind_and_dump_checkpoints(self): + """dump_messages_with_checkpoints reinserts bound checkpoints after dump.""" + cp = CheckpointData(id="cp_99") + msg = Message(role="assistant", content="hi") + msg._checkpoint_after = cp + dumped = dump_messages_with_checkpoints([msg]) + assert len(dumped) == 2 + assert dumped[0]["role"] == "assistant" + assert dumped[1]["role"] == "_checkpoint" + assert dumped[1]["content"]["id"] == "cp_99" + + def test_bind_checkpoint_messages_roundtrip(self): + """bind_checkpoint_messages binds checkpoints to prior messages.""" + history = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + {"role": "_checkpoint", "content": {"id": "cp_1"}}, + ] + messages = bind_checkpoint_messages(history) + assert len(messages) == 2 + assert messages[1]._checkpoint_after is not None + assert messages[1]._checkpoint_after.id == "cp_1" + + +class TestMessageSegments: + """Typed message segment subclasses.""" + + def test_assistant_message_segment(self): + """AssistantMessageSegment has role fixed to 'assistant'.""" + msg = AssistantMessageSegment(content="hello") + assert msg.role == "assistant" + + def test_user_message_segment(self): + """UserMessageSegment has role fixed to 'user'.""" + msg = UserMessageSegment(content="hello") + assert msg.role == "user" + + def test_system_message_segment(self): + """SystemMessageSegment has role fixed to 'system'.""" + msg = SystemMessageSegment(content="beep") + assert msg.role == "system" + + def test_tool_call_message_segment(self): + """ToolCallMessageSegment has role fixed to 'tool'.""" + msg = ToolCallMessageSegment(content="result", tool_call_id="c1") + assert msg.role == "tool" + + def test_checkpoint_message_segment(self): + """CheckpointMessageSegment has role fixed to '_checkpoint' and optional CheckpointData content.""" + cp = CheckpointData(id="cp_1") + msg = CheckpointMessageSegment(content=cp) + assert msg.role == "_checkpoint" + assert msg.content.id == "cp_1" + + +class TestContentParts: + """Content part subclasses.""" + + def test_text_part(self): + """TextPart holds text content.""" + tp = TextPart(text="Hello, world!") + assert tp.type == "text" + assert tp.text == "Hello, world!" + + def test_think_part(self): + """ThinkPart holds think content.""" + tp = ThinkPart(think="I need to think about this.") + assert tp.type == "think" + assert tp.think == "I need to think about this." + assert tp.encrypted is None + + def test_think_part_merge(self): + """merge_in_place appends think content.""" + t1 = ThinkPart(think="First ") + t2 = ThinkPart(think="Second") + assert t1.merge_in_place(t2) is True + assert t1.think == "First Second" + + def test_think_part_merge_non_think_returns_false(self): + """merge_in_place returns False when other is not a ThinkPart.""" + t1 = ThinkPart(think="A") + assert t1.merge_in_place("not a think part") is False + + def test_think_part_merge_encrypted_returns_false(self): + """merge_in_place returns False when self is encrypted.""" + t1 = ThinkPart(think="A", encrypted="sig1") + t2 = ThinkPart(think="B") + assert t1.merge_in_place(t2) is False + + def test_image_url_part(self): + """ImageURLPart holds an image URL.""" + part = ImageURLPart(image_url=ImageURLPart.ImageURL(url="http://example.com/img.jpg")) + assert part.type == "image_url" + assert part.image_url.url == "http://example.com/img.jpg" + + def test_image_url_with_id(self): + """ImageURLPart can include an id.""" + part = ImageURLPart( + image_url=ImageURLPart.ImageURL(url="http://example.com/img.jpg", id="img_1") + ) + assert part.image_url.id == "img_1" + + +class TestToolCall: + """ToolCall construction and serialization.""" + + def test_tool_call_minimal(self): + """ToolCall with id and function body.""" + tc = ToolCall(id="call_1", function=ToolCall.FunctionBody(name="f", arguments="{}")) + assert tc.id == "call_1" + assert tc.function.name == "f" + assert tc.function.arguments == "{}" + + def test_tool_call_extra_content(self): + """ToolCall with extra_content is serialized.""" + tc = ToolCall( + id="call_2", + function=ToolCall.FunctionBody(name="g", arguments='{"x": 1}'), + extra_content={"meta": "data"}, + ) + dumped = tc.model_dump() + assert dumped["extra_content"] == {"meta": "data"} + + def test_tool_call_extra_content_none_omitted(self): + """ToolCall with extra_content=None omits the field in serialization.""" + tc = ToolCall(id="call_3", function=ToolCall.FunctionBody(name="h", arguments="{}")) + dumped = tc.model_dump() + assert "extra_content" not in dumped + + def test_message_serialization_omits_tool_calls_when_none(self): + """Message.model_dump omits tool_calls when None.""" + msg = Message(role="assistant", content="hi") + dumped = msg.model_dump() + assert "tool_calls" not in dumped + + def test_message_serialization_omits_tool_call_id_when_none(self): + """Message.model_dump omits tool_call_id when None.""" + msg = Message(role="user", content="hi") + dumped = msg.model_dump() + assert "tool_call_id" not in dumped diff --git a/tests/unit/test_astrbot_config.py b/tests/unit/test_astrbot_config.py new file mode 100644 index 000000000..c920c98cd --- /dev/null +++ b/tests/unit/test_astrbot_config.py @@ -0,0 +1,323 @@ +"""Mock-based unit tests for AstrBotConfig.""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, mock_open, patch + +import pytest + +from astrbot.core.config.astrbot_config import AstrBotConfig, RateLimitStrategy + + +class TestRateLimitStrategy: + """Tests for the RateLimitStrategy enum.""" + + def test_stall_member(self): + assert RateLimitStrategy.STALL.value == "stall" + + def test_discard_member(self): + assert RateLimitStrategy.DISCARD.value == "discard" + + +class TestAstrBotConfigConstruction: + """Construction and file-existence branches.""" + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=False) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"version": 1}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_init_creates_file_when_missing( + self, mock_json_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """When the file does not exist, __init__ writes the default config.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"version": 1}, + ) + mock_json_dump.assert_called_once() + assert config["version"] == 1 + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key_a": "value_a", "key_b": 42}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_init_loads_existing_file( + self, mock_json_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """When the file exists, __init__ loads its contents.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key_a": "default", "key_b": 0}, + ) + assert config["key_a"] == "value_a" + assert config["key_b"] == 42 + # No extra dump for the missing-key case since all keys present + if hasattr(config, "first_deploy"): + assert not config.first_deploy + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=False) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + ) + def test_first_deploy_flag_set( + self, mock_file: MagicMock, mock_exists: MagicMock + ): + """first_deploy should be True when config file is created for the first time.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"version": 1}, + ) + assert hasattr(config, "first_deploy") + assert config.first_deploy is True + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"x": null}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_none_value_replaced_with_default( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """When a config value is null, it is replaced by the default.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"x": "fallback"}, + ) + assert config["x"] == "fallback" + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"a": 1, "c": 3}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_missing_keys_inserted_from_default( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """Keys present in default but missing from file are inserted.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"a": 1, "b": 2, "c": 3}, + ) + assert config["a"] == 1 + assert config["b"] == 2 # inserted from default + assert config["c"] == 3 + + +class TestAstrBotConfigOperations: + """Dot-notation access, save, and delete.""" + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key": "val"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_getattr_existing_key( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """__getattr__ returns the value for an existing key.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key": "val"}, + ) + assert config.key == "val" + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key": "val"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_getattr_missing_key_returns_none( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """__getattr__ returns None for a missing key.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key": "val"}, + ) + assert config.non_existent is None + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key": "val"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_setattr_updates_dict( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """__setattr__ stores the value in the dict.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key": "val"}, + ) + config.new_field = 99 + assert config["new_field"] == 99 + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"temp": "x"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_delattr_removes_from_dict_and_saves( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """__delattr__ removes the key from the dict and triggers save.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"temp": "x"}, + ) + del config.temp + assert "temp" not in config + mock_dump.assert_called() + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key": "val"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_delattr_missing_key_raises( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """Deleting a nonexistent key raises AttributeError.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key": "val"}, + ) + with pytest.raises(AttributeError): + del config.non_existent + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key": "val"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_save_config_writes_to_file( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """save_config serialises self to the config_path.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key": "val"}, + ) + config.new_thing = "test" + config.save_config() + # The file should be written; at this point json.dump was already called + # during __init__ as well, so at least one call after our modification + assert mock_dump.call_count >= 1 + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=True) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + read_data='{"key": "val"}', + ) + @patch( + "astrbot.core.config.astrbot_config.json.dump", + return_value=None, + ) + def test_save_config_with_replace( + self, mock_dump: MagicMock, mock_file: MagicMock, mock_exists: MagicMock + ): + """Passing replace_config updates self and writes.""" + config = AstrBotConfig( + config_path="/fake/cmd_config.json", + default_config={"key": "val"}, + ) + config.save_config(replace_config={"replacement": True}) + assert config["replacement"] is True + + +class TestAstrBotConfigSchema: + """Schema-based construction.""" + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=False) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + ) + def test_schema_object_type( + self, mock_file: MagicMock, mock_exists: MagicMock + ): + """A schema with type 'object' produces a nested dict.""" + schema = { + "nested": {"type": "object", "items": {"a": {"type": "int"}}}, + } + config = AstrBotConfig(config_path="/fake/cfg.json", schema=schema) + assert config.nested == {"a": 0} + + @patch("astrbot.core.config.astrbot_config.os.path.exists", return_value=False) + @patch( + "astrbot.core.config.astrbot_config.open", + new_callable=mock_open, + ) + def test_schema_unsupported_type_raises( + self, mock_file: MagicMock, mock_exists: MagicMock + ): + """An unsupported type in the schema raises TypeError.""" + schema = {"bad": {"type": "unknown_type"}} + with pytest.raises(TypeError, match="不受支持的配置类型"): + AstrBotConfig(config_path="/fake/cfg.json", schema=schema) + + +class TestAstrBotConfigCheckExist: + """check_exist edge cases.""" + + def test_check_exist_empty_path(self): + """Return False when config_path is falsy.""" + config = AstrBotConfig.__new__(AstrBotConfig) + object.__setattr__(config, "config_path", "") + assert config.check_exist() is False diff --git a/tests/unit/test_astrbot_config_mgr.py b/tests/unit/test_astrbot_config_mgr.py new file mode 100644 index 000000000..52d20ad3b --- /dev/null +++ b/tests/unit/test_astrbot_config_mgr.py @@ -0,0 +1,265 @@ +"""Mock-based unit tests for AstrBotConfigManager.""" + +from __future__ import annotations + +import uuid as uuid_mod +from unittest.mock import MagicMock, call, patch + +import pytest + +from astrbot.core.astrbot_config_mgr import ( + AstrBotConfigManager, + ConfInfo, + DEFAULT_CONFIG_CONF_INFO, +) + + +@pytest.fixture +def mock_default_config(): + return MagicMock() + + +@pytest.fixture +def mock_ucr(): + return MagicMock() + + +@pytest.fixture +def mock_sp(): + return MagicMock() + + +@pytest.fixture +def acm(mock_default_config, mock_ucr, mock_sp): + """Build an AstrBotConfigManager with all dependencies mocked.""" + with patch.object(AstrBotConfigManager, "_load_all_configs", return_value=None): + acm = AstrBotConfigManager(mock_default_config, mock_ucr, mock_sp) + return acm + + +class TestAstrBotConfigManagerConstruction: + """Construction and initialisation.""" + + def test_init_stores_dependencies(self, acm, mock_default_config, mock_ucr, mock_sp): + assert acm.sp is mock_sp + assert acm.ucr is mock_ucr + assert acm.confs["default"] is mock_default_config + + def test_init_sets_abconf_data_to_none(self, acm): + assert acm.abconf_data is None + + def test_default_conf_property(self, acm, mock_default_config): + assert acm.default_conf is mock_default_config + + +class TestGetConf: + """get_conf method.""" + + def test_get_conf_returns_default_when_umo_none(self, acm): + conf = acm.get_conf(None) + assert conf is acm.confs["default"] + + def test_get_conf_returns_default_when_umo_not_mapped(self, acm, mock_ucr): + mock_ucr.get_conf_id_for_umop.return_value = None + conf = acm.get_conf("test:group:123") + assert conf is acm.confs["default"] + + def test_get_conf_returns_mapped(self, acm, mock_ucr, mock_sp): + conf_id = "uuid-abc" + mock_ucr.get_conf_id_for_umop.return_value = conf_id + mock_sp.get.return_value = {conf_id: {"path": "abconf_uuid-abc.json", "name": "test"}} + mock_conf = MagicMock() + acm.confs[conf_id] = mock_conf + conf = acm.get_conf("qq:group:456") + assert conf is mock_conf + + def test_get_conf_fallback_when_mapped_not_loaded(self, acm, mock_ucr, mock_sp, mock_default_config): + conf_id = "uuid-missing" + mock_ucr.get_conf_id_for_umop.return_value = conf_id + mock_sp.get.return_value = {conf_id: {"path": "nope.json", "name": "x"}} + conf = acm.get_conf("qq:group:789") + assert conf is mock_default_config + + +class TestGetConfInfo: + """get_conf_info method.""" + + def test_get_conf_info_returns_default_when_unmapped(self, acm, mock_ucr): + mock_ucr.get_conf_id_for_umop.return_value = None + info = acm.get_conf_info("qq:group:1") + assert info["id"] == "default" + + def test_get_conf_info_returns_mapped_meta(self, acm, mock_ucr, mock_sp): + conf_id = "uuid-mapped" + mock_ucr.get_conf_id_for_umop.return_value = conf_id + mock_sp.get.return_value = {conf_id: {"path": "cfg.json", "name": "MyCfg"}} + info = acm.get_conf_info("qq:group:2") + assert info["id"] == conf_id + assert info["path"] == "cfg.json" + assert "umop" not in info + + +class TestGetConfList: + """get_conf_list method.""" + + def test_get_conf_list_includes_default(self, acm): + acm.abconf_data = {} + lst = acm.get_conf_list() + assert DEFAULT_CONFIG_CONF_INFO in lst + + def test_get_conf_list_returns_all_abconfs(self, acm, mock_sp): + mock_sp.get.return_value = { + "u1": {"path": "a.json", "name": "A"}, + "u2": {"path": "b.json", "name": "B"}, + } + acm.abconf_data = mock_sp.get.return_value + lst = acm.get_conf_list() + ids = {item["id"] for item in lst} + assert "u1" in ids + assert "u2" in ids + assert "default" in ids + + def test_get_conf_list_skips_non_dict(self, acm, mock_sp): + mock_sp.get.return_value = { + "u1": {"path": "a.json", "name": "A"}, + "u2": "not a dict", + } + acm.abconf_data = mock_sp.get.return_value + lst = acm.get_conf_list() + assert len(lst) == 2 # only u1 + default + + +class TestCreateConf: + """create_conf method.""" + + @patch("astrbot.core.astrbot_config_mgr.uuid.uuid4", return_value=uuid_mod.UUID("00000000-0000-0000-0000-000000000001")) + @patch("astrbot.core.astrbot_config_mgr.AstrBotConfig") + @patch("astrbot.core.astrbot_config_mgr.get_astrbot_config_path", return_value="/cfg") + def test_create_conf_creates_and_saves( + self, + mock_get_path, + mock_Config, + mock_uuid, + acm, + ): + mock_conf_instance = MagicMock() + mock_Config.return_value = mock_conf_instance + conf_id = acm.create_conf(config={"key": "val"}, name="myname") + mock_Config.assert_called_once() + mock_conf_instance.save_config.assert_called_once() + assert conf_id in acm.confs + assert acm.confs[conf_id] is mock_conf_instance + + +class TestDeleteConf: + """delete_conf method.""" + + def test_delete_conf_raises_on_default(self, acm): + with pytest.raises(ValueError, match="不能删除默认配置文件"): + acm.delete_conf("default") + + def test_delete_conf_returns_false_when_not_found(self, acm, mock_sp): + mock_sp.get.return_value = {} + result = acm.delete_conf("nonexistent") + assert result is False + + @patch("astrbot.core.astrbot_config_mgr.os.remove") + @patch("astrbot.core.astrbot_config_mgr.os.path.exists", return_value=True) + @patch("astrbot.core.astrbot_config_mgr.get_astrbot_config_path", return_value="/cfg") + def test_delete_conf_removes_file_and_mapping( + self, + mock_get_path, + mock_exists, + mock_remove, + acm, + mock_sp, + ): + conf_id = "uuid-to-delete" + mock_sp.get.return_value = {conf_id: {"path": "abconf_uuid-to-delete.json", "name": "x"}} + acm.abconf_data = mock_sp.get.return_value + + result = acm.delete_conf(conf_id) + assert result is True + mock_remove.assert_called_once() + mock_sp.put.assert_called() + + @patch("astrbot.core.astrbot_config_mgr.os.path.exists", return_value=False) + @patch("astrbot.core.astrbot_config_mgr.get_astrbot_config_path", return_value="/cfg") + def test_delete_conf_handles_missing_file( + self, mock_get_path, mock_exists, acm, mock_sp + ): + conf_id = "uuid-missing-file" + mock_sp.get.return_value = {conf_id: {"path": "gone.json", "name": "x"}} + acm.abconf_data = mock_sp.get.return_value + acm.confs[conf_id] = MagicMock() + + result = acm.delete_conf(conf_id) + assert result is True + assert conf_id not in acm.confs + + +class TestUpdateConfInfo: + """update_conf_info method.""" + + def test_update_raises_on_default(self, acm): + with pytest.raises(ValueError, match="不能更新"): + acm.update_conf_info("default", name="new") + + def test_update_returns_false_when_not_found(self, acm, mock_sp): + mock_sp.get.return_value = {} + result = acm.update_conf_info("nonexistent", name="new") + assert result is False + + def test_update_renames(self, acm, mock_sp): + conf_id = "uuid-rename" + mock_sp.get.return_value = {conf_id: {"path": "x.json", "name": "old"}} + acm.abconf_data = mock_sp.get.return_value + result = acm.update_conf_info(conf_id, name="new_name") + assert result is True + assert acm.abconf_data[conf_id]["name"] == "new_name" + + +class TestG: + """g (generic getter) method.""" + + def test_g_without_umo_uses_default(self, acm, mock_default_config): + mock_default_config.get.return_value = "fallback" + val = acm.g(umo=None, key="missing") + mock_default_config.get.assert_called_with("missing", None) + assert val == "fallback" + + def test_g_with_umo_uses_get_conf(self, acm): + fake_conf = MagicMock() + fake_conf.get.return_value = 42 + acm.get_conf = MagicMock(return_value=fake_conf) + val = acm.g(umo="qq:group:1", key="some.setting") + assert val == 42 + fake_conf.get.assert_called_with("some.setting", None) + + +class TestSaveConfMapping: + """_save_conf_mapping internal method.""" + + def test_save_conf_mapping_stores_and_updates_abconf_data(self, acm, mock_sp): + mock_sp.get.return_value = {} + acm._save_conf_mapping(abconf_path="new.json", abconf_id="new-id", abconf_name="display") + mock_sp.put.assert_called() + assert "new-id" in acm.abconf_data + + +class TestLoadConfMappingEdgeCases: + """_load_conf_mapping edge cases.""" + + def test_load_conf_mapping_with_invalid_umo_str(self, acm, mock_ucr): + """An invalid umo string that can't be parsed as MessageSession returns default.""" + mock_ucr.get_conf_id_for_umop.side_effect = Exception("parse error") + info = acm._load_conf_mapping("bad_format") + assert info["id"] == "default" + + def test_load_conf_mapping_checks_meta_is_dict(self, acm, mock_ucr, mock_sp): + """If abconf metadata is not a dict, returns default.""" + conf_id = "uuid-non-dict" + mock_ucr.get_conf_id_for_umop.return_value = conf_id + mock_sp.get.return_value = {conf_id: "not a dict"} + info = acm._load_conf_mapping("qq:friend:1") + assert info["id"] == "default" diff --git a/tests/unit/test_astrbot_star_main.py b/tests/unit/test_astrbot_star_main.py new file mode 100644 index 000000000..8b23d4b1b --- /dev/null +++ b/tests/unit/test_astrbot_star_main.py @@ -0,0 +1,267 @@ +"""Mock-based unit tests for AstrBot builtin star (Main).""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.api.event import AstrMessageEvent +from astrbot.api.message_components import Image, Plain +from astrbot.api.provider import LLMResponse, ProviderRequest +from astrbot.builtin_stars.astrbot.main import Main + + +@pytest.fixture +def mock_star_context(): + ctx = MagicMock() + ctx.astrbot_config_mgr = MagicMock() + return ctx + + +@pytest.fixture +def star(mock_star_context): + return Main(mock_star_context) + + +class TestMainConstruction: + """Construction and LTM initialisation.""" + + def test_init_stores_context(self, star, mock_star_context): + assert star.context is mock_star_context + + @patch("astrbot.builtin_stars.astrbot.main.LongTermMemory") + def test_init_creates_ltm(self, mock_ltm_cls, mock_star_context): + mock_ltm_instance = MagicMock() + mock_ltm_cls.return_value = mock_ltm_instance + s = Main(mock_star_context) + mock_ltm_cls.assert_called_once_with( + mock_star_context.astrbot_config_mgr, mock_star_context + ) + assert s.ltm is mock_ltm_instance + + @patch("astrbot.builtin_stars.astrbot.main.logger") + @patch("astrbot.builtin_stars.astrbot.main.LongTermMemory", side_effect=ValueError("fail")) + def test_init_handles_ltm_failure(self, mock_ltm_cls, mock_logger, mock_star_context): + s = Main(mock_star_context) + assert s.ltm is None + mock_logger.error.assert_called() + + +class TestLtmEnabled: + """ltm_enabled helper.""" + + def test_ltm_enabled_true_when_group_icl(self, star): + event = MagicMock() + event.unified_msg_origin = "qq:group:1" + star.context.get_config.return_value = { + "provider_ltm_settings": { + "group_icl_enable": True, + "active_reply": {"enable": False}, + }, + } + assert star.ltm_enabled(event) is True + + def test_ltm_enabled_true_when_active_reply(self, star): + event = MagicMock() + event.unified_msg_origin = "qq:group:1" + star.context.get_config.return_value = { + "provider_ltm_settings": { + "group_icl_enable": False, + "active_reply": {"enable": True}, + }, + } + assert star.ltm_enabled(event) is True + + def test_ltm_enabled_false_when_both_disabled(self, star): + event = MagicMock() + event.unified_msg_origin = "qq:group:1" + star.context.get_config.return_value = { + "provider_ltm_settings": { + "group_icl_enable": False, + "active_reply": {"enable": False}, + }, + } + assert star.ltm_enabled(event) is False + + +class TestOnMessage: + """on_message handler.""" + + @pytest.mark.asyncio + async def test_on_message_skips_when_no_plain_or_image(self, star): + event = MagicMock() + event.message_obj.message = [MagicMock(spec=object)] # neither Plain nor Image + gen = star.on_message(event) + items = [item async for item in gen] + assert items == [] + + @pytest.mark.asyncio + async def test_on_message_skips_when_ltm_disabled(self, star): + event = MagicMock() + event.message_obj.message = [Plain(text="hello")] + star.ltm_enabled = MagicMock(return_value=False) + gen = star.on_message(event) + items = [item async for item in gen] + assert items == [] + + @pytest.mark.asyncio + @patch.object(Main, "ltm_enabled", return_value=True) + async def test_on_message_records_context(self, mock_enabled, star): + event = MagicMock() + event.message_obj.message = [Plain(text="hello")] + star.ltm = MagicMock() + star.ltm.need_active_reply = AsyncMock(return_value=False) + star.ltm.handle_message = AsyncMock() + + star.context.get_config.return_value = { + "provider_ltm_settings": {"group_icl_enable": True, "active_reply": {"enable": False}}, + } + gen = star.on_message(event) + items = [item async for item in gen] + star.ltm.handle_message.assert_awaited_once() + + @pytest.mark.asyncio + @patch.object(Main, "ltm_enabled", return_value=True) + async def test_on_message_active_reply_triggers_llm(self, mock_enabled, star): + event = MagicMock() + event.message_obj.message = [Plain(text="hello")] + event.message_str = "hello" + event.session_id = "sess-1" + event.unified_msg_origin = "qq:group:1" + + star.ltm = MagicMock() + star.ltm.need_active_reply = AsyncMock(return_value=True) + star.ltm.handle_message = AsyncMock() + + awaitables = [] + + def request_llm(prompt, session_id, conversation): + nonlocal awaitables + awaitables.append((prompt, session_id, conversation)) + return MagicMock() + + event.request_llm = MagicMock(side_effect=request_llm) + + provider = MagicMock() + provider.text_chat = AsyncMock() + star.context.get_using_provider.return_value = provider + star.context.conversation_manager.get_curr_conversation_id = AsyncMock(return_value="cid-1") + conv = MagicMock() + star.context.conversation_manager.get_conversation = AsyncMock(return_value=conv) + + star.context.get_config.return_value = { + "provider_ltm_settings": {"group_icl_enable": True, "active_reply": {"enable": True}}, + } + + gen = star.on_message(event) + items = [item async for item in gen] + assert len(awaitables) == 1 + assert awaitables[0][0] == "hello" + + @pytest.mark.asyncio + @patch.object(Main, "ltm_enabled", return_value=True) + async def test_on_message_logs_when_no_provider(self, mock_enabled, star): + event = MagicMock() + event.message_obj.message = [Plain(text="hi")] + star.ltm = MagicMock() + star.ltm.need_active_reply = AsyncMock(return_value=True) + star.ltm.handle_message = AsyncMock() + star.context.get_using_provider.return_value = None + star.context.get_config.return_value = { + "provider_ltm_settings": {"group_icl_enable": True, "active_reply": {"enable": True}}, + } + gen = star.on_message(event) + items = [item async for item in gen] + assert items == [] + + +class TestDecorateLlmReq: + """Decorating LLM requests.""" + + @pytest.mark.asyncio + async def test_decorate_llm_req_calls_ltm(self, star): + event = MagicMock() + req = ProviderRequest(prompt="test") + star.ltm = MagicMock() + star.ltm.on_req_llm = AsyncMock() + star.ltm_enabled = MagicMock(return_value=True) + await star.decorate_llm_req(event, req) + star.ltm.on_req_llm.assert_awaited_once_with(event, req) + + @pytest.mark.asyncio + async def test_decorate_llm_req_skips_when_disabled(self, star): + event = MagicMock() + req = MagicMock() + star.ltm = MagicMock() + star.ltm_enabled = MagicMock(return_value=False) + await star.decorate_llm_req(event, req) + star.ltm.on_req_llm.assert_not_awaited() + + @pytest.mark.asyncio + async def test_decorate_llm_req_skips_when_ltm_none(self, star): + star.ltm = None + event = MagicMock() + req = MagicMock() + await star.decorate_llm_req(event, req) + + +class TestRecordLlmResp: + """Recording LLM responses.""" + + @pytest.mark.asyncio + async def test_record_llm_resp_calls_ltm(self, star): + event = MagicMock() + resp = MagicMock(spec=LLMResponse) + star.ltm = MagicMock() + star.ltm.after_req_llm = AsyncMock() + star.ltm_enabled = MagicMock(return_value=True) + await star.record_llm_resp_to_ltm(event, resp) + star.ltm.after_req_llm.assert_awaited_once_with(event, resp) + + @pytest.mark.asyncio + async def test_record_llm_resp_skips_when_disabled(self, star): + event = MagicMock() + resp = MagicMock(spec=LLMResponse) + star.ltm = MagicMock() + star.ltm_enabled = MagicMock(return_value=False) + await star.record_llm_resp_to_ltm(event, resp) + star.ltm.after_req_llm.assert_not_awaited() + + +class TestAfterMessageSent: + """After-message-sent handler.""" + + @pytest.mark.asyncio + async def test_after_message_sent_cleans_session(self, star): + event = MagicMock() + event.get_extra.return_value = True + star.ltm = MagicMock() + star.ltm.remove_session = AsyncMock(return_value=3) + star.ltm_enabled = MagicMock(return_value=True) + await star.after_message_sent(event) + event.get_extra.assert_called_with("_clean_ltm_session", False) + star.ltm.remove_session.assert_awaited_once_with(event) + + @pytest.mark.asyncio + async def test_after_message_sent_skips_when_not_flagged(self, star): + event = MagicMock() + event.get_extra.return_value = False + star.ltm = MagicMock() + star.ltm_enabled = MagicMock(return_value=True) + await star.after_message_sent(event) + star.ltm.remove_session.assert_not_awaited() + + @pytest.mark.asyncio + async def test_after_message_sent_skips_when_ltm_disabled(self, star): + event = MagicMock() + star.ltm = MagicMock() + star.ltm_enabled = MagicMock(return_value=False) + await star.after_message_sent(event) + star.ltm.remove_session.assert_not_awaited() + + @pytest.mark.asyncio + async def test_after_message_sent_skips_when_ltm_none(self, star): + star.ltm = None + event = MagicMock() + await star.after_message_sent(event) diff --git a/tests/unit/test_command_parser.py b/tests/unit/test_command_parser.py new file mode 100644 index 000000000..c0658a2c1 --- /dev/null +++ b/tests/unit/test_command_parser.py @@ -0,0 +1,97 @@ +"""Unit tests for astrbot.core.utils.command_parser. + +Covers CommandTokens data-holder and CommandParserMixin parse/regex helpers. +""" + +from astrbot.core.utils.command_parser import CommandParserMixin, CommandTokens + + +# --------------------------------------------------------------------------- +# CommandTokens +# --------------------------------------------------------------------------- + + +class TestCommandTokens: + def test_default_attributes(self): + tokens = CommandTokens() + assert tokens.tokens == [] + assert tokens.len == 0 + + def test_get_returns_none_for_out_of_bounds_negative(self): + tokens = CommandTokens() + assert tokens.get(-1) is None + + def test_get_returns_none_for_out_of_bounds_positive(self): + tokens = CommandTokens() + tokens.tokens = ["a", "b"] + tokens.len = 2 + assert tokens.get(2) is None + assert tokens.get(100) is None + + def test_get_returns_stripped_token(self): + tokens = CommandTokens() + tokens.tokens = [" hello ", "world "] + tokens.len = 2 + assert tokens.get(0) == "hello" + assert tokens.get(1) == "world" + + def test_get_from_empty_list(self): + tokens = CommandTokens() + tokens.tokens = [] + tokens.len = 0 + assert tokens.get(0) is None + + +# --------------------------------------------------------------------------- +# CommandParserMixin +# --------------------------------------------------------------------------- + + +class _ConcreteParser(CommandParserMixin): + """Minimal concrete subclass so the mixin can be instantiated.""" + + +class TestCommandParserMixin: + def setup_method(self) -> None: + self.parser = _ConcreteParser() + + def test_parse_commands_returns_command_tokens(self): + result = self.parser.parse_commands("hello world") + assert isinstance(result, CommandTokens) + assert result.len == 2 + + def test_parse_commands_splits_on_whitespace(self): + result = self.parser.parse_commands("one two\tthree\nfour") + assert result.tokens == ["one", "two", "three", "four"] + assert result.len == 4 + + def test_parse_commands_empty_string_yields_single_empty_token(self): + result = self.parser.parse_commands("") + assert result.tokens == [""] + assert result.len == 1 + + def test_parse_commands_only_whitespace(self): + result = self.parser.parse_commands(" \t \n ") + assert result.tokens == ["", ""] # \S+ splits on whitespace runs + # Actually re.split(r"\s+", " \t \n ") = ['', ''] + + def test_regex_match_returns_true_on_match(self): + assert self.parser.regex_match("hello world", "hello") is True + + def test_regex_match_returns_false_on_no_match(self): + assert self.parser.regex_match("hello world", "goodbye") is False + + def test_regex_match_multiline(self): + text = "line1\nline2\nline3" + assert self.parser.regex_match(text, "^line2$") is True + + def test_regex_match_with_digits(self): + assert self.parser.regex_match("abc123def", r"\d+") is True + + def test_regex_match_with_special_regex_chars(self): + """Special regex characters should be treated as regex, not literal.""" + assert self.parser.regex_match("price is 10.50", r"\d+\.\d+") is True + + def test_regex_match_empty_pattern(self): + """An empty pattern matches any string (re.search always finds '').""" + assert self.parser.regex_match("anything", "") is True diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py new file mode 100644 index 000000000..7e525c711 --- /dev/null +++ b/tests/unit/test_context.py @@ -0,0 +1,529 @@ +"""Unit tests for astrbot.core.star.context. + +Tests Context methods with mock-based isolation. +""" + +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest + +from astrbot.core.star.context import Context +from astrbot.core.star.star import StarMetadata, star_map, star_registry +from astrbot.core.star.star_handler import ( + EventType, + StarHandlerMetadata, + star_handlers_registry, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_dependencies(): + """Create fully isolated mocks for all Context constructor dependencies.""" + return { + "event_queue": MagicMock(), + "config": MagicMock(), + "db": MagicMock(), + "provider_manager": MagicMock(), + "platform_manager": MagicMock(), + "conversation_manager": MagicMock(), + "message_history_manager": MagicMock(), + "persona_manager": MagicMock(), + "astrbot_config_mgr": MagicMock(), + "knowledge_base_manager": MagicMock(), + "cron_manager": MagicMock(), + } + + +@pytest.fixture +def context(mock_dependencies): + """Create a Context with all mocks.""" + return Context( + mock_dependencies["event_queue"], + mock_dependencies["config"], + mock_dependencies["db"], + mock_dependencies["provider_manager"], + mock_dependencies["platform_manager"], + mock_dependencies["conversation_manager"], + mock_dependencies["message_history_manager"], + mock_dependencies["persona_manager"], + mock_dependencies["astrbot_config_mgr"], + mock_dependencies["knowledge_base_manager"], + mock_dependencies["cron_manager"], + ) + + +# --------------------------------------------------------------------------- +# Constructor / Init +# --------------------------------------------------------------------------- + + +class TestContextInit: + """Context constructor and attribute initialization.""" + + def test_init_sets_attributes(self, context, mock_dependencies): + """All constructor arguments are stored as instance attributes.""" + assert context._event_queue is mock_dependencies["event_queue"] + assert context._config is mock_dependencies["config"] + assert context._db is mock_dependencies["db"] + assert context.provider_manager is mock_dependencies["provider_manager"] + assert context.platform_manager is mock_dependencies["platform_manager"] + assert context.conversation_manager is mock_dependencies["conversation_manager"] + assert context.message_history_manager is mock_dependencies["message_history_manager"] + assert context.persona_manager is mock_dependencies["persona_manager"] + assert context.astrbot_config_mgr is mock_dependencies["astrbot_config_mgr"] + assert context.kb_manager is mock_dependencies["knowledge_base_manager"] + assert context.cron_manager is mock_dependencies["cron_manager"] + + def test_init_sets_empty_registrations(self, context): + """Runtime registration containers start empty.""" + assert context._registered_web_apis == [] + assert context._register_tasks == [] + assert context._star_manager is None + + +# --------------------------------------------------------------------------- +# get_using_provider +# --------------------------------------------------------------------------- + + +class TestGetUsingProvider: + """Context.get_using_provider() behavior.""" + + def test_returns_provider_when_found(self, context, mock_dependencies): + """get_using_provider returns the provider from provider_manager.""" + mock_provider = MagicMock() + mock_provider.__class__.__name__ = "Provider" + mock_dependencies["provider_manager"].get_using_provider.return_value = ( + mock_provider + ) + result = context.get_using_provider("test_umo") + assert result is mock_provider + mock_dependencies["provider_manager"].get_using_provider.assert_called_once() + + def test_returns_none_when_not_found(self, context, mock_dependencies): + """get_using_provider returns None when no provider is available.""" + mock_dependencies["provider_manager"].get_using_provider.return_value = None + result = context.get_using_provider("test_umo") + assert result is None + + def test_raises_value_error_on_wrong_type(self, context, mock_dependencies): + """get_using_provider raises ValueError when provider is wrong type.""" + mock_dependencies["provider_manager"].get_using_provider.return_value = ( + "not_a_provider" + ) + with pytest.raises(ValueError, match="类型不正确"): + context.get_using_provider("test_umo") + + +# --------------------------------------------------------------------------- +# get_config +# --------------------------------------------------------------------------- + + +class TestGetConfig: + """Context.get_config() behavior.""" + + def test_returns_default_config_without_umo(self, context, mock_dependencies): + """get_config() returns _config when umo is None.""" + result = context.get_config() + assert result is mock_dependencies["config"] + + def test_returns_umo_config_when_umo_provided(self, context, mock_dependencies): + """get_config(umo) delegates to astrbot_config_mgr.""" + mock_umo_config = MagicMock() + mock_dependencies["astrbot_config_mgr"].get_conf.return_value = mock_umo_config + result = context.get_config("test_umo") + assert result is mock_umo_config + mock_dependencies["astrbot_config_mgr"].get_conf.assert_called_once_with( + "test_umo" + ) + + +# --------------------------------------------------------------------------- +# get_registered_star / get_all_stars +# --------------------------------------------------------------------------- + + +class TestGetRegisteredStar: + """Context.get_registered_star() behavior.""" + + @patch("astrbot.core.star.context.star_registry", new_callable=list) + def test_finds_star_by_name(self, mock_registry, context): + """get_registered_star returns the matching StarMetadata.""" + s1 = MagicMock(spec=StarMetadata, name="plugin_a") + s1.name = "plugin_a" + s2 = MagicMock(spec=StarMetadata, name="plugin_b") + s2.name = "plugin_b" + mock_registry.extend([s1, s2]) + result = context.get_registered_star("plugin_a") + assert result is s1 + + @patch("astrbot.core.star.context.star_registry", new_callable=list) + def test_returns_none_when_not_found(self, mock_registry, context): + """get_registered_star returns None when no plugin matches.""" + mock_registry.clear() + result = context.get_registered_star("nonexistent") + assert result is None + + def test_get_all_stars_returns_registry(self, context): + """get_all_stars returns the module-level star_registry list.""" + assert context.get_all_stars() is star_registry + + +# --------------------------------------------------------------------------- +# register_commands +# --------------------------------------------------------------------------- + + +class TestRegisterCommands: + """Context.register_commands() behavior.""" + + def test_registers_command_handler(self, context): + """register_commands creates a StarHandlerMetadata and appends it.""" + async def fake_handler(): + pass + + fake_handler.__module__ = "data.plugins.test.main" + fake_handler.__qualname__ = "my_command" + context.register_commands( + star_name="test_star", + command_name="/hello", + desc="Says hello", + priority=5, + awaitable=fake_handler, + ) + + handlers = star_handlers_registry.get_handlers_by_module_name( + "data.plugins.test.main" + ) + assert len(handlers) == 1 + md = handlers[0] + assert md.event_type == EventType.AdapterMessageEvent + assert md.desc == "Says hello" + assert md.handler is fake_handler + + def test_registers_regex_command(self, context): + """register_commands with use_regex=True adds a RegexFilter.""" + async def fake_handler(): + pass + + fake_handler.__module__ = "data.plugins.test.main" + fake_handler.__qualname__ = "regex_cmd" + context.register_commands( + star_name="test_star", + command_name=r"hello.*", + desc="Regex command", + priority=1, + awaitable=fake_handler, + use_regex=True, + ) + + handlers = star_handlers_registry.get_handlers_by_module_name( + "data.plugins.test.main" + ) + assert len(handlers) == 1 + md = handlers[0] + # Should have a RegexFilter (not a CommandFilter) + from astrbot.core.star.filter.regex import RegexFilter + + assert any(isinstance(f, RegexFilter) for f in md.event_filters) + + def test_registers_command_with_ignore_prefix(self, context): + """register_commands with use_regex=False adds a CommandFilter.""" + async def fake_handler(): + pass + + fake_handler.__module__ = "data.plugins.test.main" + fake_handler.__qualname__ = "cmd_no_regex" + context.register_commands( + star_name="test_star", + command_name="/test", + desc="A command", + priority=1, + awaitable=fake_handler, + use_regex=False, + ) + + handlers = star_handlers_registry.get_handlers_by_module_name( + "data.plugins.test.main" + ) + assert len(handlers) == 1 + md = handlers[0] + from astrbot.core.star.filter.command import CommandFilter + + assert any(isinstance(f, CommandFilter) for f in md.event_filters) + + +# --------------------------------------------------------------------------- +# register_web_api +# --------------------------------------------------------------------------- + + +class TestRegisterWebApi: + """Context.register_web_api() behavior.""" + + def test_registers_new_api(self, context): + """register_web_api appends a new web API route.""" + async def handler(): + pass + + context.register_web_api( + route="/api/test", + view_handler=handler, + methods=["GET"], + desc="Test endpoint", + ) + assert len(context._registered_web_apis) == 1 + assert context._registered_web_apis[0] == ( + "/api/test", + handler, + ["GET"], + "Test endpoint", + ) + + def test_replaces_existing_route_with_same_methods(self, context): + """register_web_api replaces a previously registered API with same route and methods.""" + async def old_handler(): + pass + + async def new_handler(): + pass + + context.register_web_api( + route="/api/test", + view_handler=old_handler, + methods=["GET"], + desc="Old handler", + ) + context.register_web_api( + route="/api/test", + view_handler=new_handler, + methods=["GET"], + desc="New handler", + ) + assert len(context._registered_web_apis) == 1 + assert context._registered_web_apis[0][1] is new_handler + assert context._registered_web_apis[0][3] == "New handler" + + def test_allows_different_methods_on_same_route(self, context): + """register_web_api treats different HTTP methods as separate entries.""" + async def handler(): + pass + + context.register_web_api( + route="/api/test", + view_handler=handler, + methods=["GET"], + desc="GET handler", + ) + context.register_web_api( + route="/api/test", + view_handler=handler, + methods=["POST"], + desc="POST handler", + ) + assert len(context._registered_web_apis) == 2 + + +# --------------------------------------------------------------------------- +# add_llm_tools +# --------------------------------------------------------------------------- + + +class TestAddLLMTools: + """Context.add_llm_tools() behavior.""" + + def test_adds_tool_with_module_path(self, context, mock_dependencies): + """add_llm_tools appends tools and sets handler_module_path.""" + tool = MagicMock() + tool.name = "test_tool" + tool.__module__ = "astrbot.builtin_stars.my_plugin.main" + + mock_dependencies["provider_manager"].llm_tools.func_list = [] + context.add_llm_tools(tool) + + assert tool.handler_module_path == "astrbot.builtin_stars.my_plugin.main" + assert tool in mock_dependencies["provider_manager"].llm_tools.func_list + + def test_replaces_existing_tool_with_same_name(self, context, mock_dependencies): + """add_llm_tools replaces an existing tool with the same name.""" + old_tool = MagicMock() + old_tool.name = "dup_tool" + old_tool.__module__ = "old.module" + + new_tool = MagicMock() + new_tool.name = "dup_tool" + new_tool.__module__ = "new.module" + + mock_dependencies["provider_manager"].llm_tools.func_list = [old_tool] + context.add_llm_tools(new_tool) + + assert old_tool not in mock_dependencies["provider_manager"].llm_tools.func_list + assert new_tool in mock_dependencies["provider_manager"].llm_tools.func_list + + +# --------------------------------------------------------------------------- +# register_llm_tool (deprecated) +# --------------------------------------------------------------------------- + + +class TestRegisterLLMTool: + """Context.register_llm_tool() (deprecated) behavior.""" + + def test_registers_handler_and_adds_func(self, context, mock_dependencies): + """register_llm_tool creates a StarHandlerMetadata and adds to func_list.""" + async def handler(): + pass + + handler.__module__ = "data.plugins.p.main" + handler.__qualname__ = "my_func" + mock_dependencies["provider_manager"].llt = MagicMock() + mock_dependencies["provider_manager"].llm_tools.func_list = [] + + context.register_llm_tool( + name="my_tool", + func_args=[{"type": "string", "name": "arg1"}], + desc="My tool", + func_obj=handler, + ) + + # Check StarHandlerMetadata was added + handlers = star_handlers_registry.get_handlers_by_module_name( + "data.plugins.p.main" + ) + assert len(handlers) == 1 + assert handlers[0].event_type == EventType.OnLLMRequestEvent + + def test_calls_add_func_on_manager(self, context, mock_dependencies): + """register_llm_tool delegates to llm_tools.add_func.""" + async def handler(): + pass + + handler.__module__ = "mod" + handler.__qualname__ = "fn" + mock_dependencies["provider_manager"].llm_tools.func_list = [] + + context.register_llm_tool( + name="my_tool", + func_args=[{"type": "string", "name": "arg1"}], + desc="desc", + func_obj=handler, + ) + + mock_dependencies["provider_manager"].llm_tools.add_func.assert_called_once_with( + "my_tool", + [{"type": "string", "name": "arg1"}], + "desc", + handler, + ) + + +# --------------------------------------------------------------------------- +# register_task / reset_runtime_registrations +# --------------------------------------------------------------------------- + + +class TestRegisterTask: + """Context.register_task() behavior.""" + + def test_appends_task(self, context): + """register_task appends to _register_tasks.""" + async def task(): + pass + + context.register_task(task, "A background task") + assert task in context._register_tasks + + +class TestResetRuntimeRegistrations: + """Context.reset_runtime_registrations() behavior.""" + + def test_clears_web_apis_and_tasks(self, context): + """reset_runtime_registrations clears both containers.""" + async def handler(): + pass + + async def task(): + pass + + context.register_web_api("/api/test", handler, ["GET"], "desc") + context.register_task(task, "desc") + assert len(context._registered_web_apis) == 1 + assert len(context._register_tasks) == 1 + + context.reset_runtime_registrations() + + assert context._registered_web_apis == [] + assert context._register_tasks == [] + + +# --------------------------------------------------------------------------- +# get_platform / get_platform_inst (deprecated) +# --------------------------------------------------------------------------- + + +class TestGetPlatform: + """Context.get_platform() (deprecated) and get_platform_inst().""" + + def test_get_platform_by_string_name(self, context, mock_dependencies): + """get_platform finds a platform by its name string.""" + platform = MagicMock() + platform.meta().name = "telegram" + mock_dependencies["platform_manager"].platform_insts = [platform] + + result = context.get_platform("telegram") + assert result is platform + + def test_get_platform_returns_none_when_not_found(self, context, mock_dependencies): + """get_platform returns None when no platform matches.""" + mock_dependencies["platform_manager"].platform_insts = [] + result = context.get_platform("telegram") + assert result is None + + def test_get_platform_inst_by_id(self, context, mock_dependencies): + """get_platform_inst finds a platform by its meta id.""" + platform = MagicMock() + platform.meta().id = "test_id" + mock_dependencies["platform_manager"].platform_insts = [platform] + + result = context.get_platform_inst("test_id") + assert result is platform + + def test_get_platform_inst_returns_none(self, context, mock_dependencies): + """get_platform_inst returns None when no platform matches.""" + platform = MagicMock() + platform.meta().id = "other_id" + mock_dependencies["platform_manager"].platform_insts = [platform] + result = context.get_platform_inst("nonexistent") + assert result is None + + +# --------------------------------------------------------------------------- +# get_db +# --------------------------------------------------------------------------- + + +class TestGetDB: + """Context.get_db() behavior.""" + + def test_returns_db(self, context, mock_dependencies): + """get_db returns the _db instance.""" + assert context.get_db() is mock_dependencies["db"] + + +# --------------------------------------------------------------------------- +# get_event_queue +# --------------------------------------------------------------------------- + + +class TestGetEventQueue: + """Context.get_event_queue() behavior.""" + + def test_returns_event_queue(self, context, mock_dependencies): + """get_event_queue returns the _event_queue.""" + assert context.get_event_queue() is mock_dependencies["event_queue"] diff --git a/tests/unit/test_conversation_mgr.py b/tests/unit/test_conversation_mgr.py new file mode 100644 index 000000000..8d4e444f7 --- /dev/null +++ b/tests/unit/test_conversation_mgr.py @@ -0,0 +1,785 @@ +"""Unit tests for astrbot.core.conversation_mgr.ConversationManager.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.conversation_mgr import ConversationManager +from astrbot.core.db.po import Conversation, ConversationV2 + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_conv_v2( + conversation_id="conv-1", + platform_id="test", + user_id="test_user", + title="Test Title", + persona_id=None, + content=None, + token_usage=None, +): + """Factory for ConversationV2 with sensible defaults.""" + from datetime import datetime, timezone + + return ConversationV2( + conversation_id=conversation_id, + platform_id=platform_id, + user_id=user_id, + title=title, + persona_id=persona_id, + content=content or [], + token_usage=token_usage, + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_db(): + """Create a mock BaseDatabase.""" + db = MagicMock() + db.create_conversation = AsyncMock() + db.get_conversation_by_id = AsyncMock() + db.get_conversations = AsyncMock() + db.get_filtered_conversations = AsyncMock() + db.delete_conversation = AsyncMock() + db.delete_conversations_by_user_id = AsyncMock() + db.update_conversation = AsyncMock() + return db + + +@pytest.fixture +def mgr(mock_db): + """Create a ConversationManager with mocked database.""" + return ConversationManager(mock_db) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestConversationManagerInit: + """Tests for ConversationManager.__init__().""" + + def test_init_sets_attributes(self, mock_db): + """Verify __init__ sets all initial attributes.""" + mgr = ConversationManager(mock_db) + assert mgr.db is mock_db + assert mgr.session_conversations == {} + assert mgr.save_interval == 60 + assert mgr._on_session_deleted_callbacks == [] + + def test_register_on_session_deleted(self, mgr): + """Verify register_on_session_deleted adds callback.""" + async def callback(_umo: str): + pass + + mgr.register_on_session_deleted(callback) + assert callback in mgr._on_session_deleted_callbacks + + +class TestConversationManagerConvertV2ToV1: + """Tests for _convert_conv_from_v2_to_v1().""" + + def test_converts_v2_to_v1(self, mgr): + """Verify basic conversion of ConversationV2 to Conversation.""" + conv_v2 = _make_conv_v2( + conversation_id="c1", + platform_id="qq", + user_id="user:123", + title="Chat", + persona_id="p1", + content=[{"role": "user", "content": "hi"}], + token_usage=100, + ) + conv = mgr._convert_conv_from_v2_to_v1(conv_v2) + assert isinstance(conv, Conversation) + assert conv.cid == "c1" + assert conv.platform_id == "qq" + assert conv.user_id == "user:123" + assert conv.title == "Chat" + assert conv.persona_id == "p1" + assert conv.token_usage == 100 + assert json.loads(conv.history) == [{"role": "user", "content": "hi"}] + + def test_converts_empty_content(self, mgr): + """Verify conversion handles None content.""" + conv_v2 = _make_conv_v2(content=None) + conv = mgr._convert_conv_from_v2_to_v1(conv_v2) + assert json.loads(conv.history) == [] + + def test_converts_null_timestamps(self, mgr): + """Verify conversion handles None timestamps.""" + conv_v2 = _make_conv_v2() + conv_v2.created_at = None + conv_v2.updated_at = None + conv = mgr._convert_conv_from_v2_to_v1(conv_v2) + assert conv.created_at == 0 + assert conv.updated_at == 0 + + +class TestConversationManagerNewConversation: + """Tests for new_conversation().""" + + @pytest.mark.asyncio + async def test_new_conversation_creates_and_caches(self, mgr, mock_db): + """Verify new_conversation creates via DB and caches the ID.""" + created_conv = _make_conv_v2(conversation_id="new-conv") + mock_db.create_conversation.return_value = created_conv + + with patch("astrbot.core.conversation_mgr.sp.session_put", AsyncMock()) as sp_put: + cid = await mgr.new_conversation( + "qq:group:123", + platform_id="qq", + content=[{"role": "user", "content": "hi"}], + title="New Chat", + persona_id="p1", + ) + + assert cid == "new-conv" + assert mgr.session_conversations["qq:group:123"] == "new-conv" + mock_db.create_conversation.assert_awaited_once_with( + user_id="qq:group:123", + platform_id="qq", + content=[{"role": "user", "content": "hi"}], + title="New Chat", + persona_id="p1", + ) + sp_put.assert_awaited_once_with("qq:group:123", "sel_conv_id", "new-conv") + + @pytest.mark.asyncio + async def test_new_conversation_infers_platform_id(self, mgr, mock_db): + """Verify platform_id is inferred from unified_msg_origin when not provided.""" + created_conv = _make_conv_v2(conversation_id="cid") + mock_db.create_conversation.return_value = created_conv + + with patch("astrbot.core.conversation_mgr.sp.session_put", AsyncMock()): + cid = await mgr.new_conversation("discord:dm:456") + + assert cid == "cid" + mock_db.create_conversation.assert_awaited_once_with( + user_id="discord:dm:456", + platform_id="discord", + content=None, + title=None, + persona_id=None, + ) + + @pytest.mark.asyncio + async def test_new_conversation_fallback_platform(self, mgr, mock_db): + """Verify platform_id falls back to 'unknown' when it cannot be inferred.""" + created_conv = _make_conv_v2(conversation_id="cid") + mock_db.create_conversation.return_value = created_conv + + with patch("astrbot.core.conversation_mgr.sp.session_put", AsyncMock()): + cid = await mgr.new_conversation("short") + + assert cid == "cid" + mock_db.create_conversation.assert_awaited_once_with( + user_id="short", + platform_id="unknown", + content=None, + title=None, + persona_id=None, + ) + + +class TestConversationManagerSwitchConversation: + """Tests for switch_conversation().""" + + @pytest.mark.asyncio + async def test_switch_conversation(self, mgr): + """Verify switch updates cache and persists to session prefs.""" + mgr.session_conversations["origin:1"] = "old-id" + + with patch("astrbot.core.conversation_mgr.sp.session_put", AsyncMock()) as sp_put: + await mgr.switch_conversation("origin:1", "new-id") + + assert mgr.session_conversations["origin:1"] == "new-id" + sp_put.assert_awaited_once_with("origin:1", "sel_conv_id", "new-id") + + +class TestConversationManagerDeleteConversation: + """Tests for delete_conversation().""" + + @pytest.mark.asyncio + async def test_delete_current_conversation(self, mgr, mock_db): + """Verify deleting the current conversation clears cache.""" + mgr.session_conversations["origin:1"] = "conv-id" + mock_db.get_conversation_by_id = AsyncMock(return_value=None) + mock_db.delete_conversation = AsyncMock() + + with ( + patch( + "astrbot.core.conversation_mgr.sp.session_remove", + AsyncMock(), + ) as sp_remove, + patch.object( + mgr, + "get_curr_conversation_id", + AsyncMock(return_value="conv-id"), + ), + ): + await mgr.delete_conversation("origin:1") + + mock_db.delete_conversation.assert_awaited_once_with(cid="conv-id") + assert "origin:1" not in mgr.session_conversations + sp_remove.assert_awaited_once_with("origin:1", "sel_conv_id") + + @pytest.mark.asyncio + async def test_delete_specific_conversation(self, mgr, mock_db): + """Verify deleting a non-current conversation by ID.""" + mgr.session_conversations["origin:1"] = "current-id" + mock_db.delete_conversation = AsyncMock() + + with ( + patch( + "astrbot.core.conversation_mgr.sp.session_remove", + AsyncMock(), + ) as sp_remove, + patch.object( + mgr, + "get_curr_conversation_id", + AsyncMock(return_value="current-id"), + ), + ): + await mgr.delete_conversation("origin:1", conversation_id="other-id") + + mock_db.delete_conversation.assert_awaited_once_with(cid="other-id") + # Current ID should NOT be removed since it differs + assert mgr.session_conversations.get("origin:1") == "current-id" + sp_remove.assert_not_called() + + @pytest.mark.asyncio + async def test_delete_no_conversation_id_fallback(self, mgr, mock_db): + """Verify when no conv ID is given and none cached, do nothing.""" + mgr.session_conversations = {} + mock_db.delete_conversation = AsyncMock() + + with patch( + "astrbot.core.conversation_mgr.sp.session_remove", + AsyncMock(), + ): + await mgr.delete_conversation("origin:1") + + mock_db.delete_conversation.assert_not_called() + + +class TestConversationManagerDeleteConversationsByUserId: + """Tests for delete_conversations_by_user_id().""" + + @pytest.mark.asyncio + async def test_delete_all_and_trigger_callbacks(self, mgr, mock_db): + """Verify deleting all conversations cleans cache and triggers callbacks.""" + mgr.session_conversations["origin:1"] = "c1" + mock_db.delete_conversations_by_user_id = AsyncMock() + + callback = AsyncMock() + mgr.register_on_session_deleted(callback) + + with patch( + "astrbot.core.conversation_mgr.sp.session_remove", + AsyncMock(), + ) as sp_remove: + await mgr.delete_conversations_by_user_id("origin:1") + + mock_db.delete_conversations_by_user_id.assert_awaited_once_with( + user_id="origin:1", + ) + assert "origin:1" not in mgr.session_conversations + sp_remove.assert_awaited_once_with("origin:1", "sel_conv_id") + callback.assert_awaited_once_with("origin:1") + + +class TestConversationManagerGetCurrConversationId: + """Tests for get_curr_conversation_id().""" + + @pytest.mark.asyncio + async def test_returns_cached_value(self, mgr): + """Verify returns cached conversation ID without hitting session prefs.""" + mgr.session_conversations["origin:1"] = "cached-id" + + with patch("astrbot.core.conversation_mgr.sp.session_get", AsyncMock()) as sp_get: + cid = await mgr.get_curr_conversation_id("origin:1") + + assert cid == "cached-id" + sp_get.assert_not_called() + + @pytest.mark.asyncio + async def test_fetches_from_session_prefs_and_caches(self, mgr): + """Verify fetches from session prefs when not cached, then caches it.""" + mgr.session_conversations = {} + + with patch( + "astrbot.core.conversation_mgr.sp.session_get", + AsyncMock(return_value="pref-id"), + ) as sp_get: + cid = await mgr.get_curr_conversation_id("origin:1") + + assert cid == "pref-id" + assert mgr.session_conversations["origin:1"] == "pref-id" + + @pytest.mark.asyncio + async def test_returns_none_when_not_found(self, mgr): + """Verify returns None when no conversation is known.""" + mgr.session_conversations = {} + + with patch( + "astrbot.core.conversation_mgr.sp.session_get", + AsyncMock(return_value=None), + ): + cid = await mgr.get_curr_conversation_id("origin:1") + + assert cid is None + + +class TestConversationManagerGetConversation: + """Tests for get_conversation().""" + + @pytest.mark.asyncio + async def test_get_existing_conversation(self, mgr, mock_db): + """Verify retrieves and converts existing conversation.""" + conv_v2 = _make_conv_v2(conversation_id="c1") + mock_db.get_conversation_by_id.return_value = conv_v2 + + conv = await mgr.get_conversation("origin:1", "c1") + + assert conv is not None + assert conv.cid == "c1" + mock_db.get_conversation_by_id.assert_awaited_once_with(cid="c1") + + @pytest.mark.asyncio + async def test_get_non_existing_not_created(self, mgr, mock_db): + """Verify returns None when not found and create_if_not_exists is False.""" + mock_db.get_conversation_by_id.return_value = None + + conv = await mgr.get_conversation("origin:1", "nonexistent") + + assert conv is None + + @pytest.mark.asyncio + async def test_get_non_existing_creates_new(self, mgr, mock_db): + """Verify creates new conversation when not found and create_if_not_exists is True.""" + mock_db.get_conversation_by_id.side_effect = [ + None, + _make_conv_v2(conversation_id="new-c1"), + ] + + with patch.object( + mgr, + "new_conversation", + AsyncMock(return_value="new-c1"), + ) as mock_new: + conv = await mgr.get_conversation( + "origin:1", + "nonexistent", + create_if_not_exists=True, + ) + + assert conv is not None + assert conv.cid == "new-c1" + mock_new.assert_awaited_once_with("origin:1") + + +class TestConversationManagerGetConversations: + """Tests for get_conversations().""" + + @pytest.mark.asyncio + async def test_get_conversations(self, mgr, mock_db): + """Verify retrieving multiple conversations.""" + convs_v2 = [ + _make_conv_v2(conversation_id="c1"), + _make_conv_v2(conversation_id="c2"), + ] + mock_db.get_conversations.return_value = convs_v2 + + convs = await mgr.get_conversations(unified_msg_origin="origin:1") + + assert len(convs) == 2 + assert convs[0].cid == "c1" + assert convs[1].cid == "c2" + mock_db.get_conversations.assert_awaited_once_with( + user_id="origin:1", + platform_id=None, + ) + + +class TestConversationManagerGetFilteredConversations: + """Tests for get_filtered_conversations().""" + + @pytest.mark.asyncio + async def test_get_filtered_conversations(self, mgr, mock_db): + """Verify filtered conversation retrieval.""" + convs_v2 = [_make_conv_v2(conversation_id="c1")] + mock_db.get_filtered_conversations.return_value = (convs_v2, 1) + + convs, cnt = await mgr.get_filtered_conversations( + page=1, + page_size=20, + platform_ids=["qq"], + search_query="test", + ) + + assert len(convs) == 1 + assert convs[0].cid == "c1" + assert cnt == 1 + mock_db.get_filtered_conversations.assert_awaited_once_with( + page=1, + page_size=20, + platform_ids=["qq"], + search_query="test", + ) + + +class TestConversationManagerUpdateConversation: + """Tests for update_conversation().""" + + @pytest.mark.asyncio + async def test_update_without_id_uses_current(self, mgr, mock_db): + """Verify update uses current conversation ID when not provided.""" + with patch.object( + mgr, + "get_curr_conversation_id", + AsyncMock(return_value="current-id"), + ): + await mgr.update_conversation( + "origin:1", + history=[{"role": "user", "content": "hi"}], + title="Updated", + persona_id="p2", + token_usage=50, + ) + + mock_db.update_conversation.assert_awaited_once_with( + cid="current-id", + title="Updated", + persona_id="p2", + clear_persona=False, + content=[{"role": "user", "content": "hi"}], + token_usage=50, + ) + + @pytest.mark.asyncio + async def test_update_with_id(self, mgr, mock_db): + """Verify update with explicit conversation ID.""" + await mgr.update_conversation( + "origin:1", + conversation_id="explicit-id", + title="New Title", + ) + + mock_db.update_conversation.assert_awaited_once_with( + cid="explicit-id", + title="New Title", + persona_id=None, + clear_persona=False, + content=None, + token_usage=None, + ) + + @pytest.mark.asyncio + async def test_update_without_id_no_current_does_nothing(self, mgr, mock_db): + """Verify update does nothing when no ID is available.""" + with patch.object( + mgr, + "get_curr_conversation_id", + AsyncMock(return_value=None), + ): + await mgr.update_conversation("origin:1", title="New Title") + + mock_db.update_conversation.assert_not_called() + + +class TestConversationManagerUpdateConversationTitle: + """Tests for update_conversation_title().""" + + @pytest.mark.asyncio + async def test_update_title_delegates(self, mgr): + """Verify update_conversation_title delegates to update_conversation.""" + with patch.object(mgr, "update_conversation", AsyncMock()) as mock_update: + await mgr.update_conversation_title("origin:1", "New Title", "conv-id") + + mock_update.assert_awaited_once_with( + unified_msg_origin="origin:1", + conversation_id="conv-id", + title="New Title", + ) + + +class TestConversationManagerUpdateConversationPersonaId: + """Tests for update_conversation_persona_id().""" + + @pytest.mark.asyncio + async def test_update_persona_delegates(self, mgr): + """Verify update_conversation_persona_id delegates to update_conversation.""" + with patch.object(mgr, "update_conversation", AsyncMock()) as mock_update: + await mgr.update_conversation_persona_id( + "origin:1", + "persona-123", + "conv-id", + ) + + mock_update.assert_awaited_once_with( + unified_msg_origin="origin:1", + conversation_id="conv-id", + persona_id="persona-123", + ) + + +class TestConversationManagerUnsetConversationPersona: + """Tests for unset_conversation_persona().""" + + @pytest.mark.asyncio + async def test_unset_persona_delegates_with_clear(self, mgr): + """Verify unset_conversation_persona delegates with clear_persona=True.""" + with patch.object(mgr, "update_conversation", AsyncMock()) as mock_update: + await mgr.unset_conversation_persona("origin:1", "conv-id") + + mock_update.assert_awaited_once_with( + unified_msg_origin="origin:1", + conversation_id="conv-id", + clear_persona=True, + ) + + +class TestConversationManagerAddMessagePair: + """Tests for add_message_pair().""" + + @pytest.mark.asyncio + async def test_add_message_pair_dicts(self, mgr, mock_db): + """Verify adding message pair using plain dicts.""" + conv_v2 = _make_conv_v2(content=[]) + mock_db.get_conversation_by_id.return_value = conv_v2 + + await mgr.add_message_pair( + "c1", + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + ) + + mock_db.update_conversation.assert_awaited_once_with( + cid="c1", + content=[ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + ], + ) + + @pytest.mark.asyncio + async def test_add_message_pair_appends_to_existing(self, mgr, mock_db): + """Verify adding message pair appends to existing history.""" + conv_v2 = _make_conv_v2(content=[{"role": "user", "content": "prev"}]) + mock_db.get_conversation_by_id.return_value = conv_v2 + + await mgr.add_message_pair( + "c1", + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ) + + mock_db.update_conversation.assert_awaited_once_with( + cid="c1", + content=[ + {"role": "user", "content": "prev"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ], + ) + + @pytest.mark.asyncio + async def test_add_message_pair_conv_not_found(self, mgr, mock_db): + """Verify raises when conversation is not found.""" + mock_db.get_conversation_by_id.return_value = None + + with pytest.raises(Exception, match="Conversation with id nonexistent not found"): + await mgr.add_message_pair( + "nonexistent", + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ) + + @pytest.mark.asyncio + async def test_add_message_pair_segment_objects(self, mgr, mock_db): + """Verify adding message pair using UserMessageSegment / AssistantMessageSegment.""" + from astrbot.core.agent.message import ( + AssistantMessageSegment, + UserMessageSegment, + ) + + conv_v2 = _make_conv_v2(content=[]) + mock_db.get_conversation_by_id.return_value = conv_v2 + + user_msg = UserMessageSegment(content="hello") + assistant_msg = AssistantMessageSegment(content="world") + + with ( + patch.object(user_msg, "model_dump", return_value={"role": "user", "content": "hello"}), + patch.object( + assistant_msg, + "model_dump", + return_value={"role": "assistant", "content": "world"}, + ), + ): + await mgr.add_message_pair("c1", user_msg, assistant_msg) + + mock_db.update_conversation.assert_awaited_once_with( + cid="c1", + content=[ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ], + ) + + +class TestConversationManagerGetHumanReadableContext: + """Tests for get_human_readable_context().""" + + @pytest.mark.asyncio + async def test_get_context(self, mgr): + """Verify basic context formatting.""" + mgr._convert_conv_from_v2_to_v1 = MagicMock( + return_value=Conversation( + platform_id="qq", + user_id="user:1", + cid="c1", + history=json.dumps([ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + {"role": "user", "content": "how are you"}, + {"role": "assistant", "content": "fine"}, + ]), + title="Chat", + created_at=1000, + updated_at=1001, + ), + ) + + contexts, total_pages = await mgr.get_human_readable_context( + "origin:1", + "c1", + page=1, + page_size=10, + ) + + # Order: most recent pair first + assert contexts == [ + "User: how are you", + "Assistant: fine", + "User: hello", + "Assistant: world", + ] + assert total_pages == 1 + + @pytest.mark.asyncio + async def test_get_context_no_conversation(self, mgr): + """Verify returns empty when conversation not found.""" + with patch.object( + mgr, + "get_conversation", + AsyncMock(return_value=None), + ): + contexts, total_pages = await mgr.get_human_readable_context( + "origin:1", + "nonexistent", + ) + + assert contexts == [] + assert total_pages == 0 + + @pytest.mark.asyncio + async def test_get_context_pagination(self, mgr): + """Verify context pagination.""" + records = [] + for i in range(5): + records.append({"role": "user", "content": f"q{i}"}) + records.append({"role": "assistant", "content": f"a{i}"}) + + mgr._convert_conv_from_v2_to_v1 = MagicMock( + return_value=Conversation( + platform_id="qq", + user_id="user:1", + cid="c1", + history=json.dumps(records), + title="Chat", + created_at=1000, + updated_at=1001, + ), + ) + + contexts, total_pages = await mgr.get_human_readable_context( + "origin:1", + "c1", + page=1, + page_size=2, + ) + + # 5 pairs = 10 records, reversed, page_size=2 + assert len(contexts) == 2 + assert total_pages == 5 + + @pytest.mark.asyncio + async def test_get_context_with_tool_calls(self, mgr): + """Verify context handles tool_calls in assistant messages.""" + mgr._convert_conv_from_v2_to_v1 = MagicMock( + return_value=Conversation( + platform_id="qq", + user_id="user:1", + cid="c1", + history=json.dumps([ + {"role": "user", "content": "search weather"}, + { + "role": "assistant", + "tool_calls": [{"function": {"name": "get_weather"}}], + }, + ]), + title="Chat", + created_at=1000, + updated_at=1001, + ), + ) + + contexts, _ = await mgr.get_human_readable_context("origin:1", "c1") + + assert "Assistant: [函数调用]" in contexts[0] + + +class TestConversationManagerTriggerSessionDeleted: + """Tests for _trigger_session_deleted().""" + + @pytest.mark.asyncio + async def test_triggers_callbacks(self, mgr): + """Verify all callbacks are triggered.""" + cb1 = AsyncMock() + cb2 = AsyncMock() + mgr.register_on_session_deleted(cb1) + mgr.register_on_session_deleted(cb2) + + await mgr._trigger_session_deleted("origin:1") + + cb1.assert_awaited_once_with("origin:1") + cb2.assert_awaited_once_with("origin:1") + + @pytest.mark.asyncio + async def test_callback_error_does_not_block_others(self, mgr): + """Verify one failing callback does not prevent others from running.""" + cb1 = AsyncMock(side_effect=RuntimeError("fail")) + cb2 = AsyncMock() + mgr.register_on_session_deleted(cb1) + mgr.register_on_session_deleted(cb2) + + # Should not raise + await mgr._trigger_session_deleted("origin:1") + + cb2.assert_awaited_once_with("origin:1") diff --git a/tests/unit/test_cron_events.py b/tests/unit/test_cron_events.py new file mode 100644 index 000000000..db07b429f --- /dev/null +++ b/tests/unit/test_cron_events.py @@ -0,0 +1,242 @@ +"""Tests for CronMessageEvent.""" + +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.cron.events import CronMessageEvent +from astrbot.core.message.components import Plain +from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.platform.message_session import MessageSession +from astrbot.core.platform.message_type import MessageType + + +@pytest.fixture +def mock_context(): + ctx = MagicMock() + ctx.send_message = AsyncMock() + return ctx + + +@pytest.fixture +def mock_session(): + session = MagicMock(spec=MessageSession) + session.session_id = "test-session-id" + session.platform_id = "test-platform" + return session + + +class TestCronMessageEventInit: + """Tests for CronMessageEvent construction.""" + + def test_init_default_params(self, mock_context, mock_session): + """Default params produce a synthetic event with correct attributes.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Hello from cron", + ) + assert event.message_str == "Hello from cron" + assert event.is_at_or_wake_command is True + assert event.is_wake is True + assert event.session == mock_session + assert event.context_obj == mock_context + + def test_init_sender_defaults(self, mock_context, mock_session): + """Default sender_id and sender_name are 'astrbot' and 'Scheduler'.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + assert event.message_obj.self_id == "astrbot" + assert event.message_obj.sender.nickname == "Scheduler" + assert event.message_obj.sender.user_id == "test-session-id" + + def test_init_custom_sender(self, mock_context, mock_session): + """Custom sender_id and sender_name are reflected on the message object.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + sender_id="custom-bot", + sender_name="CustomName", + ) + assert event.message_obj.self_id == "custom-bot" + assert event.message_obj.sender.nickname == "CustomName" + + def test_init_with_extras(self, mock_context, mock_session): + """Extras dict is merged into _extras.""" + extras = {"origin": "api", "cron_payload": {"key": "value"}} + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + extras=extras, + ) + assert event._extras.get("origin") == "api" + assert event._extras["cron_payload"] == {"key": "value"} + + def test_init_group_message_type(self, mock_context, mock_session): + """MessageType.GROUP_MESSAGE is preserved in the message object.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + message_type=MessageType.GROUP_MESSAGE, + ) + assert event.message_obj.type == MessageType.GROUP_MESSAGE + + def test_init_platform_meta(self, mock_context, mock_session): + """Platform metadata is set to cron / CronJob / platform_id.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + assert event.platform_meta.name == "cron" + assert event.platform_meta.description == "CronJob" + assert event.platform_meta.id == "test-platform" + + def test_init_message_components(self, mock_context, mock_session): + """The message is wrapped in a Plain component.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Cron message text", + ) + assert len(event.message_obj.message) == 1 + assert isinstance(event.message_obj.message[0], Plain) + assert event.message_obj.message[0].text == "Cron message text" + assert event.message_obj.message_str == "Cron message text" + assert event.message_obj.raw_message == "Cron message text" + + def test_init_timestamp_within_range(self, mock_context, mock_session): + """Timestamp is set to the current time.""" + before = int(time.time()) + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + after = int(time.time()) + assert before <= event.message_obj.timestamp <= after + + def test_init_message_id_is_uuid_hex(self, mock_context, mock_session): + """Message ID is a 32-char hex string (uuid4 hex).""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + mid = event.message_obj.message_id + assert isinstance(mid, str) + assert len(mid) == 32 + # All characters should be valid hex digits + int(mid, 16) + + def test_init_none_extras_does_not_raise(self, mock_context, mock_session): + """Passing extras=None does not raise.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + extras=None, + ) + # Should not raise; _extras remains default + assert isinstance(event._extras, dict) + + def test_init_uses_session_session_id(self, mock_context, mock_session): + """Session ID is taken from session.session_id.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + assert event.session_id == "test-session-id" + + +class TestCronMessageEventSend: + """Tests for CronMessageEvent.send.""" + + @pytest.mark.asyncio + async def test_send_calls_context(self, mock_context, mock_session): + """send delegates to context.send_message on the original session.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + chain = MessageChain([Plain("response")]) + + with patch.object(AstrMessageEvent, "send", new_callable=AsyncMock): + await event.send(chain) + + mock_context.send_message.assert_awaited_once() + + @pytest.mark.asyncio + async def test_send_none_is_noop(self, mock_context, mock_session): + """send(None) returns immediately without calling context.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + with patch.object(AstrMessageEvent, "send", new_callable=AsyncMock) as mock_super: + await event.send(None) + + mock_context.send_message.assert_not_called() + mock_super.assert_not_called() + + @pytest.mark.asyncio + async def test_send_calls_super(self, mock_context, mock_session): + """send also delegates to super().send().""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + chain = MessageChain([Plain("response")]) + + with patch.object(AstrMessageEvent, "send", new_callable=AsyncMock) as mock_super: + await event.send(chain) + + mock_super.assert_awaited_once_with(chain) + + @pytest.mark.asyncio + async def test_send_streaming_iterates_generator(self, mock_context, mock_session): + """send_streaming calls send for each yielded chain.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + + async def gen(): + yield MessageChain([Plain("part1")]) + yield MessageChain([Plain("part2")]) + + with patch.object(event, "send", new_callable=AsyncMock) as mock_send: + await event.send_streaming(gen()) + + assert mock_send.call_count == 2 + + @pytest.mark.asyncio + async def test_send_streaming_empty_generator(self, mock_context, mock_session): + """send_streaming does not call send when the generator is empty.""" + event = CronMessageEvent( + context=mock_context, + session=mock_session, + message="Test", + ) + + async def empty_gen(): + if False: + yield # pragma: no cover + + with patch.object(event, "send", new_callable=AsyncMock) as mock_send: + await event.send_streaming(empty_gen()) + + mock_send.assert_not_called() diff --git a/tests/unit/test_cron_manager_edge.py b/tests/unit/test_cron_manager_edge.py new file mode 100644 index 000000000..6a74918c3 --- /dev/null +++ b/tests/unit/test_cron_manager_edge.py @@ -0,0 +1,428 @@ +"""Supplementary edge-case tests for CronJobManager. + +Covers validation branches, interval-based jobs, run-once auto-cleanup, +and failure paths not covered by the main test suite. +""" + +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.cron.manager import CronJobManager, CronJobSchedulingError +from astrbot.core.db.po import CronJob + + +# ---- Fixtures (self-contained) ---- + + +@pytest.fixture +def mock_db(): + db = MagicMock() + db.create_cron_job = AsyncMock() + db.get_cron_job = AsyncMock() + db.update_cron_job = AsyncMock() + db.delete_cron_job = AsyncMock() + db.list_cron_jobs = AsyncMock(return_value=[]) + return db + + +@pytest.fixture +def mock_context(): + ctx = MagicMock() + ctx.get_config = MagicMock(return_value={"admins_id": []}) + ctx.conversation_manager = MagicMock() + return ctx + + +@pytest.fixture +def cron_manager(mock_db): + return CronJobManager(mock_db) + + +# ---- add_basic_job validation edge cases ---- + + +class TestAddBasicJobEdgeCases: + """Validation and interval-based variants for add_basic_job.""" + + @pytest.mark.asyncio + async def test_add_basic_job_with_interval_seconds(self, cron_manager, mock_db): + """interval_seconds can be used instead of cron_expression.""" + cron_manager._started = True + job = CronJob(job_id="interval-job", name="Interval", job_type="basic") + mock_db.create_cron_job.return_value = job + + handler = MagicMock() + result = await cron_manager.add_basic_job( + name="Interval", + interval_seconds=300, + handler=handler, + description="Every 5 min", + ) + + assert result == job + mock_db.create_cron_job.assert_called_once() + # interval should be packed into the payload + call_payload = mock_db.create_cron_job.call_args.kwargs["payload"] + assert "interval_seconds" in call_payload + assert call_payload["interval_seconds"] == 300 + + @pytest.mark.asyncio + async def test_add_basic_job_both_cron_and_interval_raises(self, cron_manager, mock_db): + """Providing both cron_expression and interval_seconds raises ValueError.""" + handler = MagicMock() + with pytest.raises(ValueError, match="must have exactly one value"): + await cron_manager.add_basic_job( + name="Bad", + cron_expression="0 9 * * *", + interval_seconds=300, + handler=handler, + ) + + @pytest.mark.asyncio + async def test_add_basic_job_neither_cron_nor_interval_raises(self, cron_manager, mock_db): + """Providing neither cron_expression nor interval_seconds raises ValueError.""" + handler = MagicMock() + with pytest.raises(ValueError, match="must have exactly one value"): + await cron_manager.add_basic_job( + name="Bad", + handler=handler, + ) + + @pytest.mark.asyncio + async def test_add_basic_job_payload_passed_through(self, cron_manager, mock_db): + """User-supplied payload is forwarded to create_cron_job.""" + cron_manager._started = True + job = CronJob(job_id="payload-job", name="Payload", job_type="basic") + mock_db.create_cron_job.return_value = job + + handler = MagicMock() + user_payload = {"custom_key": "custom_value"} + await cron_manager.add_basic_job( + name="Payload", + cron_expression="0 9 * * *", + handler=handler, + payload=user_payload, + ) + + call_payload = mock_db.create_cron_job.call_args.kwargs["payload"] + assert call_payload["custom_key"] == "custom_value" + + @pytest.mark.asyncio + async def test_add_basic_job_non_persistent_not_scheduled(self, cron_manager, mock_db): + """A disabled non-persistent job is stored but not scheduled.""" + job = CronJob( + job_id="np-job", name="NonPersist", job_type="basic", enabled=False + ) + mock_db.create_cron_job.return_value = job + + handler = MagicMock() + with patch.object(cron_manager, "_schedule_job") as mock_schedule: + result = await cron_manager.add_basic_job( + name="NonPersist", + cron_expression="0 9 * * *", + handler=handler, + persistent=False, + enabled=False, + ) + + assert result == job + mock_schedule.assert_not_called() + + +# ---- _run_job full lifecycle edge cases ---- + + +class TestRunJobEdgeCases: + """Full lifecycle for _run_job covering status transitions and cleanup.""" + + @pytest.mark.asyncio + async def test_run_job_basic_completed(self, cron_manager, mock_db): + """A basic job transitions to 'completed' after successful run.""" + job = CronJob( + job_id="basic-done", + name="Done", + job_type="basic", + enabled=True, + cron_expression="0 9 * * *", + ) + mock_db.get_cron_job.return_value = job + + handler = MagicMock(return_value=None) + cron_manager._basic_handlers["basic-done"] = handler + + await cron_manager._run_job("basic-done") + + # Status should be updated twice: running -> completed + update_calls = [ + c for c in mock_db.update_cron_job.call_args_list + ] + # At minimum one call with status="running" and one with status="completed" + running_calls = [ + c for c in update_calls + if c.kwargs.get("status") == "running" + ] + completed_calls = [ + c for c in update_calls + if c.kwargs.get("status") == "completed" + ] + assert len(running_calls) >= 1, "Expected a 'running' status update" + assert len(completed_calls) >= 1, "Expected a 'completed' status update" + handler.assert_called_once() + + @pytest.mark.asyncio + async def test_run_job_basic_failed(self, cron_manager, mock_db): + """A basic job transitions to 'failed' and records the error.""" + mock_db.get_cron_job.return_value = CronJob( + job_id="basic-fail", + name="Fail", + job_type="basic", + enabled=True, + cron_expression="0 9 * * *", + ) + # Remove the handler so _run_basic_job raises RuntimeError + if "basic-fail" in cron_manager._basic_handlers: + del cron_manager._basic_handlers["basic-fail"] + + await cron_manager._run_job("basic-fail") + + # final update should have status="failed" and last_error set + last_call = mock_db.update_cron_job.call_args_list[-1] + assert last_call.kwargs["status"] == "failed" + assert last_call.kwargs["last_error"] is not None + + @pytest.mark.asyncio + async def test_run_job_run_once_deletes_after_completion(self, cron_manager, mock_db): + """A run_once job is deleted after it completes successfully.""" + job = CronJob( + job_id="once-job", + name="Once", + job_type="basic", + enabled=True, + run_once=True, + cron_expression="0 9 * * *", + ) + mock_db.get_cron_job.return_value = job + + handler = MagicMock(return_value=None) + cron_manager._basic_handlers["once-job"] = handler + + with patch.object(cron_manager, "delete_job") as mock_delete: + await cron_manager._run_job("once-job") + + mock_delete.assert_awaited_once_with("once-job") + + @pytest.mark.asyncio + async def test_run_job_active_agent_raises_on_missing_session(self, cron_manager, mock_db): + """run_job on an active_agent job without session payload raises.""" + job = CronJob( + job_id="aa-no-session", + name="NoSession", + job_type="active_agent", + enabled=True, + cron_expression="0 9 * * *", + payload={}, # No "session" key + ) + mock_db.get_cron_job.return_value = job + + await cron_manager._run_job("aa-no-session") + + # Should not crash the manager; error is caught and logged, + # status should be "failed" + last_call = mock_db.update_cron_job.call_args_list[-1] + assert last_call.kwargs["status"] == "failed" + assert "missing session" in (last_call.kwargs.get("last_error") or "").lower() + + +# ---- update_job and sync_from_db ---- + + +class TestUpdateAndSyncEdgeCases: + """Edge cases for update_job and sync_from_db.""" + + @pytest.mark.asyncio + async def test_update_job_reschedules_when_enabled(self, cron_manager, mock_db): + """update_job re-schedules a job that was previously disabled.""" + job_id = "re-enable-job" + updated_job = CronJob( + job_id=job_id, + name="ReEnabled", + job_type="basic", + cron_expression="0 9 * * *", + enabled=True, # Now enabled + ) + mock_db.update_cron_job.return_value = updated_job + + with patch.object(cron_manager, "_schedule_job") as mock_schedule: + result = await cron_manager.update_job(job_id, enabled=True) + + assert result == updated_job + mock_schedule.assert_called_once_with(updated_job) + + @pytest.mark.asyncio + async def test_update_job_removes_scheduled_when_disabled(self, cron_manager, mock_db): + """update_job removes the job from the scheduler when disabled.""" + job_id = "disable-job" + updated_job = CronJob( + job_id=job_id, + name="Disabled", + job_type="basic", + cron_expression="0 9 * * *", + enabled=False, # Now disabled, should not schedule + ) + mock_db.update_cron_job.return_value = updated_job + + with patch.object(cron_manager, "_remove_scheduled") as mock_remove: + with patch.object(cron_manager, "_schedule_job") as mock_schedule: + result = await cron_manager.update_job(job_id, enabled=False) + + assert result == updated_job + mock_remove.assert_called_once_with(job_id) + mock_schedule.assert_not_called() + + @pytest.mark.asyncio + async def test_sync_from_db_schedules_basic_with_handler(self, cron_manager, mock_db): + """sync_from_db schedules basic jobs when their handler is registered.""" + job = CronJob( + job_id="sync-basic", + name="SyncBasic", + job_type="basic", + cron_expression="0 9 * * *", + enabled=True, + persistent=True, + ) + mock_db.list_cron_jobs.return_value = [job] + cron_manager._basic_handlers["sync-basic"] = MagicMock() + + with patch.object(cron_manager, "_schedule_job") as mock_schedule: + await cron_manager.sync_from_db() + + mock_schedule.assert_called_once_with(job) + + @pytest.mark.asyncio + async def test_sync_from_db_skips_basic_without_handler(self, cron_manager, mock_db): + """sync_from_db skips basic jobs that have no registered handler.""" + job = CronJob( + job_id="orphan-basic", + name="Orphan", + job_type="basic", + cron_expression="0 9 * * *", + enabled=True, + persistent=True, + ) + mock_db.list_cron_jobs.return_value = [job] + + with patch.object(cron_manager, "_schedule_job") as mock_schedule: + await cron_manager.sync_from_db() + + mock_schedule.assert_not_called() + + @pytest.mark.asyncio + async def test_sync_from_db_schedules_active_agent_without_handler(self, cron_manager, mock_db): + """Active-agent jobs are scheduled regardless of handler registration.""" + job = CronJob( + job_id="sync-active", + name="SyncActive", + job_type="active_agent", + cron_expression="0 9 * * *", + enabled=True, + persistent=True, + ) + mock_db.list_cron_jobs.return_value = [job] + + with patch.object(cron_manager, "_schedule_job") as mock_schedule: + await cron_manager.sync_from_db() + + mock_schedule.assert_called_once_with(job) + + +# ---- _schedule_job trigger variants ---- + + +class TestScheduleJobTriggers: + """Trigger type selection in _schedule_job.""" + + @pytest.mark.asyncio + async def test_schedule_interval_trigger(self, cron_manager, mock_context): + """_schedule_job creates an IntervalTrigger when payload has interval_seconds.""" + job = CronJob( + job_id="interval-trigger", + name="Interval", + job_type="basic", + cron_expression="0 9 * * *", + enabled=True, + payload={"interval_seconds": 600}, + ) + mock_db = cron_manager.db + mock_db.list_cron_jobs = AsyncMock(return_value=[]) + mock_db.update_cron_job = AsyncMock() + await cron_manager.start(mock_context) + + cron_manager._schedule_job(job) + + aps_job = cron_manager.scheduler.get_job("interval-trigger") + assert aps_job is not None + # The trigger should be an IntervalTrigger + from apscheduler.triggers.interval import IntervalTrigger + + assert isinstance(aps_job.trigger, IntervalTrigger) + assert aps_job.trigger.interval.total_seconds() == 600 + + @pytest.mark.asyncio + async def test_schedule_invalid_cron_raises_scheduling_error(self, cron_manager, mock_context): + """An invalid cron expression raises CronJobSchedulingError.""" + job = CronJob( + job_id="bad-cron", + name="BadCron", + job_type="basic", + cron_expression="not-a-valid-cron", + enabled=True, + ) + mock_db = cron_manager.db + mock_db.list_cron_jobs = AsyncMock(return_value=[]) + mock_db.update_cron_job = AsyncMock() + await cron_manager.start(mock_context) + + with pytest.raises(CronJobSchedulingError): + cron_manager._schedule_job(job) + + @pytest.mark.asyncio + async def test_schedule_run_once_without_run_at_raises(self, cron_manager, mock_context): + """A run_once job without run_at in payload or expression raises CronJobSchedulingError.""" + job = CronJob( + job_id="no-run-at", + name="NoRunAt", + job_type="active_agent", + cron_expression=None, + enabled=True, + run_once=True, + payload={}, + ) + mock_db = cron_manager.db + mock_db.list_cron_jobs = AsyncMock(return_value=[]) + mock_db.update_cron_job = AsyncMock() + await cron_manager.start(mock_context) + + with pytest.raises(CronJobSchedulingError): + cron_manager._schedule_job(job) + + @pytest.mark.asyncio + async def test_schedule_auto_starts_when_not_started(self, cron_manager): + """_schedule_job auto-starts the scheduler if _started is False.""" + assert cron_manager._started is False + + job = CronJob( + job_id="auto-start", + name="AutoStart", + job_type="basic", + cron_expression="0 9 * * *", + enabled=True, + ) + cron_manager.db.list_cron_jobs = AsyncMock(return_value=[]) + cron_manager.db.update_cron_job = AsyncMock() + + cron_manager._schedule_job(job) + + assert cron_manager._started is True + assert cron_manager.scheduler.get_job("auto-start") is not None diff --git a/tests/unit/test_db_init.py b/tests/unit/test_db_init.py new file mode 100644 index 000000000..1917ef69f --- /dev/null +++ b/tests/unit/test_db_init.py @@ -0,0 +1,254 @@ +"""Tests for BaseDatabase abstract interface and initialization. + +Verifies the abstract method contract, engine creation, session factory +setup, and the ``get_db`` async context manager. +""" + +import inspect +from abc import ABC +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.db import BaseDatabase + + +class TestBaseDatabaseAbstract: + """Tests for BaseDatabase's abstract nature and method count.""" + + def test_is_abstract_class(self): + """BaseDatabase uses ABCMeta and cannot be directly instantiated.""" + assert isinstance(BaseDatabase, ABC) + with pytest.raises(TypeError, match="abstract"): + BaseDatabase() + + def test_abstract_method_count(self): + """BaseDatabase defines a known number of abstract methods.""" + abs_methods = BaseDatabase.__abstractmethods__ + # The interface is large -- expect >= 90 abstract methods + assert len(abs_methods) >= 90 + # Key abstract methods should be present + assert "create_cron_job" in abs_methods + assert "update_cron_job" in abs_methods + assert "delete_cron_job" in abs_methods + assert "get_cron_job" in abs_methods + assert "list_cron_jobs" in abs_methods + assert "get_conversations" in abs_methods + assert "create_conversation" in abs_methods + assert "insert_platform_stats" in abs_methods + assert "insert_persona" in abs_methods + assert "insert_api_key" in abs_methods # kept as create_api_key is the actual name + # Deprecated methods are still abstract + assert "get_base_stats" in abs_methods + assert "get_total_message_count" in abs_methods + + def test_abstract_methods_have_docstrings_or_impl(self): + """All abstract methods should have ... or docstrings (no syntax errors at module level).""" + abs_methods = BaseDatabase.__abstractmethods__ + for name in abs_methods: + method = getattr(BaseDatabase, name) + # Should not raise AttributeError or similar + assert callable(method) + + def test_initialize_is_not_abstract(self): + """initialize() has a concrete default implementation.""" + assert "initialize" not in BaseDatabase.__abstractmethods__ + + def test_get_db_is_not_abstract(self): + """get_db() has a concrete implementation (async context manager).""" + assert "get_db" not in BaseDatabase.__abstractmethods__ + + +class TestBaseDatabaseInit: + """Tests for BaseDatabase.__init__ and engine setup.""" + + def test_init_sets_inited_false(self): + """inited starts as False.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine") as mock_engine, + patch("astrbot.core.db.async_sessionmaker") as mock_smaker, + ): + mock_engine.return_value = MagicMock() + mock_smaker.return_value = MagicMock() + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() + + assert db.inited is False + + def test_init_sqlite_adds_connect_args_timeout(self): + """SQLite URL adds timeout=30 to connect_args.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine") as mock_engine, + patch("astrbot.core.db.async_sessionmaker") as mock_smaker, + ): + mock_engine.return_value = MagicMock() + mock_smaker.return_value = MagicMock() + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() # noqa: F841 + + mock_engine.assert_called_once() + call_kwargs = mock_engine.call_args.kwargs + assert call_kwargs["connect_args"] == {"timeout": 30} + + def test_init_non_sqlite_does_not_add_connect_args(self): + """Non-SQLite URL omits connect_args (no timeout).""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine") as mock_engine, + patch("astrbot.core.db.async_sessionmaker") as mock_smaker, + ): + mock_engine.return_value = MagicMock() + mock_smaker.return_value = MagicMock() + BaseDatabase.DATABASE_URL = "postgresql+asyncpg://localhost/db" + db = BaseDatabase() # noqa: F841 + + call_kwargs = mock_engine.call_args.kwargs + assert "connect_args" not in call_kwargs or call_kwargs["connect_args"] == {} + + def test_init_creates_async_session_local(self): + """A session factory is created from the engine.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine") as mock_engine, + patch("astrbot.core.db.async_sessionmaker") as mock_smaker, + ): + mock_engine.return_value = MagicMock() + mock_smaker.return_value = MagicMock() + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() # noqa: F841 + + mock_smaker.assert_called_once() + # The session maker is bound to the engine with expire_on_commit=False + assert mock_smaker.call_args.kwargs.get("expire_on_commit") is False + + def test_init_engine_echo_false_future_true(self): + """Engine is created with echo=False and future=True.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine") as mock_engine, + patch("astrbot.core.db.async_sessionmaker") as mock_smaker, + ): + mock_engine.return_value = MagicMock() + mock_smaker.return_value = MagicMock() + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() # noqa: F841 + + call_kwargs = mock_engine.call_args.kwargs + assert call_kwargs["echo"] is False + assert call_kwargs["future"] is True + + def test_same_url_passed_to_engine(self): + """DATABASE_URL is passed to create_async_engine.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine") as mock_engine, + patch("astrbot.core.db.async_sessionmaker") as mock_smaker, + ): + mock_engine.return_value = MagicMock() + mock_smaker.return_value = MagicMock() + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///custom.db" + db = BaseDatabase() # noqa: F841 + + assert mock_engine.call_args[0][0] == "sqlite+aiosqlite:///custom.db" + + +class TestBaseDatabaseInitialize: + """Tests for the concrete initialize method.""" + + @pytest.mark.asyncio + async def test_initialize_does_not_raise(self): + """initialize() is concrete and can be called without error.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine"), + patch("astrbot.core.db.async_sessionmaker"), + ): + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() + + await db.initialize() # should not raise + + +class TestBaseDatabaseGetDb: + """Tests for the get_db async context manager.""" + + @pytest.mark.asyncio + async def test_get_db_calls_initialize_when_not_inited(self): + """get_db calls initialize() if inited is False.""" + from sqlalchemy.ext.asyncio import AsyncSession + + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine"), + patch("astrbot.core.db.async_sessionmaker"), + ): + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() + + db.inited = False + + # Mock the session factory + mock_session = MagicMock(spec=AsyncSession) + mock_cm = AsyncMock() + mock_cm.__aenter__ = AsyncMock(return_value=mock_session) + mock_cm.__aexit__ = AsyncMock(return_value=None) + db.AsyncSessionLocal = MagicMock(return_value=mock_cm) + + with patch.object(db, "initialize", new_callable=AsyncMock) as mock_init: + async with db.get_db() as session: + assert session == mock_session + + mock_init.assert_awaited_once() + + @pytest.mark.asyncio + async def test_get_db_does_not_call_initialize_when_inited(self): + """get_db skips initialize() if inited is already True.""" + from sqlalchemy.ext.asyncio import AsyncSession + + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine"), + patch("astrbot.core.db.async_sessionmaker"), + ): + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() + + db.inited = True + + mock_session = MagicMock(spec=AsyncSession) + mock_cm = AsyncMock() + mock_cm.__aenter__ = AsyncMock(return_value=mock_session) + mock_cm.__aexit__ = AsyncMock(return_value=None) + db.AsyncSessionLocal = MagicMock(return_value=mock_cm) + + with patch.object(db, "initialize", new_callable=AsyncMock) as mock_init: + async with db.get_db() as session: + assert session == mock_session + + mock_init.assert_not_called() + + @pytest.mark.asyncio + async def test_get_db_sets_inited_after_initialize(self): + """inited is set to True after initialize completes.""" + with ( + patch.object(BaseDatabase, "__abstractmethods__", frozenset()), + patch("astrbot.core.db.create_async_engine"), + patch("astrbot.core.db.async_sessionmaker"), + ): + BaseDatabase.DATABASE_URL = "sqlite+aiosqlite:///test.db" + db = BaseDatabase() + + db.inited = False + + mock_session = MagicMock() + mock_cm = AsyncMock() + mock_cm.__aenter__ = AsyncMock(return_value=mock_session) + mock_cm.__aexit__ = AsyncMock(return_value=None) + db.AsyncSessionLocal = MagicMock(return_value=mock_cm) + + async with db.get_db(): + pass + + assert db.inited is True diff --git a/tests/unit/test_db_po.py b/tests/unit/test_db_po.py new file mode 100644 index 000000000..a0c489cdc --- /dev/null +++ b/tests/unit/test_db_po.py @@ -0,0 +1,454 @@ +"""Tests for database PO (Persistent Object) model classes. + +These tests verify construction, field defaults, and type correctness +for all model classes defined in ``astrbot.core.db.po``. +""" + +from datetime import datetime, timezone + +import pytest + +from astrbot.core.db.po import ( + ApiKey, + Attachment, + ChatUIProject, + CommandConfig, + CommandConflict, + Conversation, + ConversationV2, + CronJob, + Persona, + PersonaFolder, + Personality, + PlatformMessageHistory, + PlatformSession, + PlatformStat, + Preference, + ProviderStat, + SessionProjectRelation, + Stats, + TimestampMixin, + WebChatThread, +) + + +class TestCronJob: + """Tests for the CronJob SQLModel.""" + + def test_minimal_construction(self): + """CronJob can be created with only required fields.""" + job = CronJob(name="test-job", job_type="basic") + assert job.name == "test-job" + assert job.job_type == "basic" + assert job.enabled is True + assert job.persistent is True + assert job.run_once is False + assert job.status == "scheduled" + assert job.cron_expression is None + assert job.job_id is not None # auto-generated via default_factory + + def test_full_construction(self): + """CronJob accepts all optional fields.""" + job = CronJob( + name="full-job", + job_type="active_agent", + cron_expression="0 9 * * *", + timezone="UTC", + payload={"session": "test:group:1"}, + description="A full job", + enabled=False, + persistent=True, + run_once=True, + status="running", + ) + assert job.name == "full-job" + assert job.job_type == "active_agent" + assert job.cron_expression == "0 9 * * *" + assert job.timezone == "UTC" + assert job.payload == {"session": "test:group:1"} + assert job.description == "A full job" + assert job.enabled is False + assert job.run_once is True + assert job.status == "running" + + def test_auto_timestamps(self): + """CronJob inherits TimestampMixin which auto-generates created_at/updated_at.""" + job = CronJob(name="ts-test", job_type="basic") + assert isinstance(job.created_at, datetime) + assert isinstance(job.updated_at, datetime) + + def test_job_id_auto_generated(self): + """Each instance gets a unique job_id.""" + job1 = CronJob(name="a", job_type="basic") + job2 = CronJob(name="b", job_type="basic") + assert job1.job_id != job2.job_id + + +class TestConversationV2: + """Tests for the ConversationV2 SQLModel.""" + + def test_minimal_construction(self): + """ConversationV2 can be created with required fields.""" + conv = ConversationV2(platform_id="qq", user_id="user-1") + assert conv.platform_id == "qq" + assert conv.user_id == "user-1" + assert conv.conversation_id is not None + assert conv.title is None + assert conv.persona_id is None + assert conv.token_usage == 0 + + def test_full_construction(self): + """ConversationV2 accepts all optional fields.""" + conv = ConversationV2( + platform_id="webchat", + user_id="admin", + content=[{"role": "user", "content": "hi"}], + title="Chat", + persona_id="p1", + token_usage=42, + ) + assert conv.content == [{"role": "user", "content": "hi"}] + assert conv.title == "Chat" + assert conv.persona_id == "p1" + assert conv.token_usage == 42 + + +class TestConversationDataclass: + """Tests for the deprecated Conversation dataclass.""" + + def test_construction(self): + """Conversation dataclass sets all fields.""" + conv = Conversation( + platform_id="qq", + user_id="u1", + cid="abc-123", + history="[{}]", + title="My Chat", + persona_id="p1", + created_at=1000, + updated_at=2000, + token_usage=50, + ) + assert conv.platform_id == "qq" + assert conv.user_id == "u1" + assert conv.cid == "abc-123" + assert conv.history == "[{}]" + assert conv.title == "My Chat" + assert conv.persona_id == "p1" + assert conv.created_at == 1000 + assert conv.updated_at == 2000 + assert conv.token_usage == 50 + + def test_defaults(self): + """Conversation dataclass has sensible defaults.""" + conv = Conversation(platform_id="qq", user_id="u1", cid="x") + assert conv.history == "" + assert conv.title == "" + assert conv.persona_id == "" + assert conv.created_at == 0 + assert conv.token_usage == 0 + + +class TestPersona: + """Tests for the Persona SQLModel.""" + + def test_minimal_construction(self): + """Persona with only required fields.""" + p = Persona(persona_id="helper", system_prompt="You are helpful.") + assert p.persona_id == "helper" + assert p.system_prompt == "You are helpful." + assert p.begin_dialogs is None + assert p.tools is None + assert p.skills is None + assert p.custom_error_message is None + assert p.sort_order == 0 + + def test_with_all_fields(self): + """Persona with all optional fields.""" + p = Persona( + persona_id="custom", + system_prompt="Be concise.", + begin_dialogs=["Hello"], + tools=["search", "calc"], + skills=["weather"], + custom_error_message="Sorry, try later.", + folder_id="folder-1", + sort_order=5, + ) + assert p.begin_dialogs == ["Hello"] + assert p.tools == ["search", "calc"] + assert p.skills == ["weather"] + assert p.custom_error_message == "Sorry, try later." + assert p.folder_id == "folder-1" + assert p.sort_order == 5 + + +class TestPersonaFolder: + """Tests for the PersonaFolder SQLModel.""" + + def test_construction(self): + """PersonaFolder creation with required fields.""" + folder = PersonaFolder(name="My Folder") + assert folder.name == "My Folder" + assert folder.folder_id is not None + assert folder.parent_id is None + assert folder.description is None + assert folder.sort_order == 0 + + def test_nested_folder(self): + """PersonaFolder can have a parent_id.""" + child = PersonaFolder(name="Child", parent_id="parent-uuid", description="Nested") + assert child.parent_id == "parent-uuid" + assert child.description == "Nested" + + +class TestApiKey: + """Tests for the ApiKey SQLModel.""" + + def test_construction(self): + """ApiKey with required fields.""" + key = ApiKey( + name="dev-key", + key_hash="sha256:abc123", + key_prefix="astr_", + created_by="admin", + ) + assert key.name == "dev-key" + assert key.key_hash == "sha256:abc123" + assert key.key_prefix == "astr_" + assert key.created_by == "admin" + assert key.scopes is None + assert key.last_used_at is None + assert key.expires_at is None + assert key.revoked_at is None + assert key.key_id is not None + + def test_with_scopes(self): + """ApiKey can have scopes and expiry.""" + expires = datetime.now(timezone.utc) + key = ApiKey( + name="scoped-key", + key_hash="sha256:xyz", + key_prefix="astr_", + created_by="admin", + scopes=["read", "write"], + expires_at=expires, + ) + assert key.scopes == ["read", "write"] + assert key.expires_at == expires + + +class TestPlatformStat: + """Tests for the PlatformStat SQLModel.""" + + def test_construction(self): + """PlatformStat with all fields.""" + ts = datetime.now(timezone.utc) + stat = PlatformStat( + timestamp=ts, + platform_id="qq_bot", + platform_type="aiocqhttp", + count=5, + ) + assert stat.timestamp == ts + assert stat.platform_id == "qq_bot" + assert stat.platform_type == "aiocqhttp" + assert stat.count == 5 + + def test_default_count(self): + """PlatformStat defaults count to 0.""" + stat = PlatformStat( + timestamp=datetime.now(timezone.utc), + platform_id="test", + platform_type="test", + ) + assert stat.count == 0 + + +class TestProviderStat: + """Tests for the ProviderStat SQLModel.""" + + def test_construction(self): + """ProviderStat with required fields.""" + ps = ProviderStat(umo="test:private:1", provider_id="openai") + assert ps.umo == "test:private:1" + assert ps.provider_id == "openai" + assert ps.agent_type == "internal" + assert ps.status == "completed" + assert ps.token_input_other == 0 + assert ps.token_input_cached == 0 + assert ps.token_output == 0 + + def test_with_stats(self): + """ProviderStat records token and timing stats.""" + ps = ProviderStat( + umo="test:private:1", + provider_id="anthropic", + provider_model="claude-3", + conversation_id="conv-1", + status="completed", + agent_type="cron", + token_input_other=100, + token_input_cached=50, + token_output=200, + start_time=1000.0, + end_time=1005.0, + time_to_first_token=0.5, + ) + assert ps.provider_model == "claude-3" + assert ps.conversation_id == "conv-1" + assert ps.token_input_other == 100 + assert ps.token_input_cached == 50 + assert ps.token_output == 200 + assert ps.start_time == 1000.0 + assert ps.end_time == 1005.0 + assert ps.time_to_first_token == 0.5 + + +class TestPreference: + """Tests for the Preference SQLModel.""" + + def test_construction(self): + """Preference with required fields.""" + pref = Preference( + scope="plugin", + scope_id="my_plugin", + key="theme", + value={"color": "dark"}, + ) + assert pref.scope == "plugin" + assert pref.scope_id == "my_plugin" + assert pref.key == "theme" + assert pref.value == {"color": "dark"} + + +class TestOtherModels: + """Tests for remaining SQLModel classes.""" + + def test_attachment(self): + """Attachment model construction.""" + att = Attachment(path="/tmp/file.png", type="image", mime_type="image/png") + assert att.path == "/tmp/file.png" + assert att.type == "image" + assert att.mime_type == "image/png" + assert att.attachment_id is not None + + def test_webchat_thread(self): + """WebChatThread model construction.""" + thread = WebChatThread( + creator="user-1", + parent_session_id="session-1", + parent_message_id=42, + base_checkpoint_id="ckpt-1", + selected_text="selected text", + ) + assert thread.creator == "user-1" + assert thread.parent_session_id == "session-1" + assert thread.parent_message_id == 42 + assert thread.base_checkpoint_id == "ckpt-1" + assert thread.selected_text == "selected text" + assert thread.thread_id is not None + + def test_platform_session(self): + """PlatformSession model construction.""" + sess = PlatformSession( + creator="user-1", + platform_id="webchat", + display_name="My Chat", + ) + assert sess.creator == "user-1" + assert sess.platform_id == "webchat" + assert sess.display_name == "My Chat" + assert sess.is_group == 0 + assert sess.session_id is not None + + def test_chatui_project(self): + """ChatUIProject model construction.""" + proj = ChatUIProject( + creator="admin", + title="My Project", + emoji="star", + description="A test project", + ) + assert proj.creator == "admin" + assert proj.title == "My Project" + assert proj.emoji == "star" + assert proj.description == "A test project" + assert proj.project_id is not None + + def test_session_project_relation(self): + """SessionProjectRelation model construction.""" + rel = SessionProjectRelation(session_id="sess-1", project_id="proj-1") + assert rel.session_id == "sess-1" + assert rel.project_id == "proj-1" + + def test_command_config(self): + """CommandConfig model construction.""" + cc = CommandConfig( + handler_full_name="plugin.command", + plugin_name="test-plugin", + module_path="plugins.test", + original_command="/test", + ) + assert cc.handler_full_name == "plugin.command" + assert cc.plugin_name == "test-plugin" + assert cc.module_path == "plugins.test" + assert cc.original_command == "/test" + assert cc.enabled is True + assert cc.auto_managed is False + + def test_command_conflict(self): + """CommandConflict model construction.""" + cf = CommandConflict( + conflict_key="/greet", + handler_full_name="p1.greet", + plugin_name="plugin1", + ) + assert cf.conflict_key == "/greet" + assert cf.handler_full_name == "p1.greet" + assert cf.plugin_name == "plugin1" + assert cf.status == "pending" + assert cf.auto_generated is False + + def test_platform_message_history(self): + """PlatformMessageHistory model construction.""" + pmh = PlatformMessageHistory( + platform_id="qq", + user_id="user-1", + content={"text": "hello"}, + ) + assert pmh.platform_id == "qq" + assert pmh.user_id == "user-1" + assert pmh.content == {"text": "hello"} + assert pmh.sender_id is None + assert pmh.sender_name is None + assert pmh.llm_checkpoint_id is None + + def test_timestamp_mixin_fields(self): + """TimestampMixin provides created_at and updated_at.""" + ts = datetime.now(timezone.utc) + mixin = TimestampMixin() + # created_at and updated_at have default factories + assert isinstance(mixin.created_at, datetime) + assert isinstance(mixin.updated_at, datetime) + + def test_personality_typeddict(self): + """Personality TypedDict can be constructed with all keys.""" + personality: Personality = { + "prompt": "You are a bot.", + "name": "Bot", + "begin_dialogs": ["Hello"], + "mood_imitation_dialogs": ["Hi there"], + "tools": ["search"], + "skills": ["weather"], + "custom_error_message": "Oops", + "_begin_dialogs_processed": [{"role": "user", "content": "Hello"}], + "_mood_imitation_dialogs_processed": "Hi there", + } + assert personality["prompt"] == "You are a bot." + assert personality["name"] == "Bot" + + def test_stats_dataclass(self): + """Stats dataclass holds a list of Platforms.""" + stats = Stats() + assert stats.platform == [] diff --git a/tests/unit/test_default_config.py b/tests/unit/test_default_config.py new file mode 100644 index 000000000..388ec1b2a --- /dev/null +++ b/tests/unit/test_default_config.py @@ -0,0 +1,211 @@ +"""Mock-based unit tests for default config constants and structure.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from astrbot.core.config.default import ( + DB_PATH, + DEFAULT_CONFIG, + DEFAULT_VALUE_MAP, + VERSION, + WEBHOOK_SUPPORTED_PLATFORMS, +) + + +class TestVersionConstant: + """Tests for the VERSION constant.""" + + def test_version_is_string(self): + assert isinstance(VERSION, str) + + def test_version_format(self): + parts = VERSION.split(".") + assert len(parts) == 3 + for p in parts: + assert p.isdigit() + + +class TestDBPath: + """Tests for DB_PATH.""" + + @patch("astrbot.core.config.default.get_astrbot_data_path", return_value="/data") + def test_db_path_uses_data_path(self, mock_get_path): + """DB_PATH should join data path with the database filename.""" + # Reimport to pick up the patched value + import importlib + + from astrbot.core.config import default as default_mod + + importlib.reload(default_mod) + assert "data_v4.db" in default_mod.DB_PATH + + def test_db_path_ends_with_db(self): + assert DB_PATH.endswith(".db") + + +class TestDEFAULT_VALUE_MAP: + """Tests for DEFAULT_VALUE_MAP structure.""" + + def test_contains_expected_types(self): + expected = {"int", "float", "bool", "string", "text", "list", "file", "object", "template_list"} + assert set(DEFAULT_VALUE_MAP.keys()) == expected + + def test_default_values_are_correct_types(self): + assert isinstance(DEFAULT_VALUE_MAP["int"], int) + assert isinstance(DEFAULT_VALUE_MAP["float"], float) + assert isinstance(DEFAULT_VALUE_MAP["bool"], bool) + assert isinstance(DEFAULT_VALUE_MAP["string"], str) + assert isinstance(DEFAULT_VALUE_MAP["text"], str) + assert isinstance(DEFAULT_VALUE_MAP["list"], list) + assert isinstance(DEFAULT_VALUE_MAP["file"], list) + assert isinstance(DEFAULT_VALUE_MAP["object"], dict) + assert isinstance(DEFAULT_VALUE_MAP["template_list"], list) + + def test_specific_default_values(self): + assert DEFAULT_VALUE_MAP["int"] == 0 + assert DEFAULT_VALUE_MAP["float"] == 0.0 + assert DEFAULT_VALUE_MAP["bool"] is False + assert DEFAULT_VALUE_MAP["string"] == "" + assert DEFAULT_VALUE_MAP["list"] == [] + assert DEFAULT_VALUE_MAP["object"] == {} + assert DEFAULT_VALUE_MAP["template_list"] == [] + + +class TestWEBHOOK_SUPPORTED_PLATFORMS: + """Tests for the webhook platforms list.""" + + def test_is_list_of_strings(self): + assert isinstance(WEBHOOK_SUPPORTED_PLATFORMS, list) + for p in WEBHOOK_SUPPORTED_PLATFORMS: + assert isinstance(p, str) + + def test_contains_key_platforms(self): + assert "qq_official_webhook" in WEBHOOK_SUPPORTED_PLATFORMS + assert "weixin_official_account" in WEBHOOK_SUPPORTED_PLATFORMS + assert "slack" in WEBHOOK_SUPPORTED_PLATFORMS + assert "lark" in WEBHOOK_SUPPORTED_PLATFORMS + + +class TestDEFAULT_CONFIGStructure: + """Tests for the top-level keys and nested structure of DEFAULT_CONFIG.""" + + def test_top_level_keys(self): + expected_keys = { + "config_version", + "platform_settings", + "provider_sources", + "provider", + "provider_settings", + "subagent_orchestrator", + "provider_stt_settings", + "provider_tts_settings", + "provider_ltm_settings", + "content_safety", + "admins_id", + "t2i", + "http_proxy", + "no_proxy", + "dashboard", + "platform", + "platform_specific", + "wake_prefix", + "log_level", + "persona", + "timezone", + "callback_api_base", + "default_kb_collection", + "plugin_set", + "kb_names", + "kb_fusion_top_k", + "kb_final_top_k", + "kb_agentic_mode", + "disable_builtin_commands", + "t2i_word_threshold", + "t2i_strategy", + "t2i_endpoint", + "t2i_use_file_service", + "t2i_active_template", + "log_file_enable", + "log_file_path", + "log_file_max_mb", + "temp_dir_max_size", + "trace_enable", + "trace_log_enable", + "trace_log_path", + "trace_log_max_mb", + "pip_install_arg", + "pypi_index_url", + } + for key in expected_keys: + assert key in DEFAULT_CONFIG, f"Missing top-level key: {key}" + + def test_config_version_is_two(self): + assert DEFAULT_CONFIG["config_version"] == 2 + + def test_platform_settings_structure(self): + ps = DEFAULT_CONFIG["platform_settings"] + assert "unique_session" in ps + assert "rate_limit" in ps + assert "reply_prefix" in ps + assert isinstance(ps["rate_limit"], dict) + assert ps["rate_limit"]["strategy"] in ("stall", "discard") + + def test_provider_ltm_settings_structure(self): + ltm = DEFAULT_CONFIG["provider_ltm_settings"] + assert "group_icl_enable" in ltm + assert "group_message_max_cnt" in ltm + assert "image_caption" in ltm + assert "active_reply" in ltm + assert isinstance(ltm["active_reply"], dict) + assert ltm["active_reply"]["method"] == "possibility_reply" + + def test_dashboard_settings(self): + dash = DEFAULT_CONFIG["dashboard"] + assert dash["enable"] is True + assert dash["username"] == "astrbot" + assert dash["port"] == 6185 + assert "ssl" in dash + + def test_provider_settings_defaults(self): + ps = DEFAULT_CONFIG["provider_settings"] + assert ps["enable"] is True + assert ps["default_provider_id"] == "" + assert ps["agent_runner_type"] == "local" + assert ps["llm_safety_mode"] is True + + def test_admins_id_default(self): + assert DEFAULT_CONFIG["admins_id"] == ["astrbot"] + + def test_wake_prefix_default(self): + assert DEFAULT_CONFIG["wake_prefix"] == ["/"] + + def test_log_level_default(self): + assert DEFAULT_CONFIG["log_level"] == "INFO" + + def test_platform_specific_contains_lark_telegram_discord(self): + ps = DEFAULT_CONFIG["platform_specific"] + assert "lark" in ps + assert "telegram" in ps + assert "discord" in ps + + def test_no_proxy_contains_localhost(self): + assert "localhost" in DEFAULT_CONFIG["no_proxy"] + + def test_subagent_orchestrator_defaults(self): + sa = DEFAULT_CONFIG["subagent_orchestrator"] + assert sa["main_enable"] is False + assert isinstance(sa["agents"], list) + + def test_quoted_message_parser_defaults(self): + qmp = DEFAULT_CONFIG["provider_settings"]["quoted_message_parser"] + assert qmp["max_component_chain_depth"] == 4 + assert qmp["max_forward_node_depth"] == 6 + assert qmp["max_forward_fetch"] == 32 + + def test_sandbox_defaults(self): + sb = DEFAULT_CONFIG["provider_settings"]["sandbox"] + assert sb["booter"] == "shipyard_neo" + assert sb["shipyard_neo_ttl"] == 3600 diff --git a/tests/unit/test_event_bus_edge.py b/tests/unit/test_event_bus_edge.py new file mode 100644 index 000000000..0526bbbf5 --- /dev/null +++ b/tests/unit/test_event_bus_edge.py @@ -0,0 +1,396 @@ +"""Supplementary edge-case tests for EventBus. + +Covers exception resilience in the dispatch loop, edge values in +_print_event, empty pipeline mappings, and rapid event delivery. +""" + +import asyncio +from contextlib import suppress +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.event_bus import EventBus + + +# ---- Fixtures ---- + + +@pytest.fixture +def event_queue(): + return asyncio.Queue() + + +@pytest.fixture +def mock_scheduler(): + scheduler = MagicMock() + scheduler.execute = AsyncMock() + return scheduler + + +@pytest.fixture +def mock_config_manager(): + config_mgr = MagicMock() + config_mgr.get_conf_info = MagicMock( + return_value={"id": "test-conf-id", "name": "Test Config"} + ) + return config_mgr + + +# ---- Exception resilience ---- + + +class TestDispatchExceptionResilience: + """The dispatch loop must survive scheduler exceptions.""" + + @pytest.mark.asyncio + async def test_exception_in_scheduler_does_not_crash_loop( + self, event_queue, mock_config_manager + ): + """After a scheduler that raises, subsequent events are still processed.""" + processed_second = asyncio.Event() + + # Scheduler 1 raises on every call + scheduler1 = MagicMock() + scheduler1.execute = AsyncMock(side_effect=RuntimeError("Boom")) + + # Scheduler 2 works normally + scheduler2 = MagicMock() + scheduler2.execute = AsyncMock() + + async def execute_second(event): # noqa: ARG001 + processed_second.set() + + scheduler2.execute.side_effect = execute_second + + def get_conf_info(origin): + if "event-1" in origin: + return {"id": "conf-1", "name": "C1"} + return {"id": "conf-2", "name": "C2"} + + mock_config_manager.get_conf_info.side_effect = get_conf_info + + mapping = {"conf-1": scheduler1, "conf-2": scheduler2} + event_bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping=mapping, + astrbot_config_mgr=mock_config_manager, + ) + + # Event 1 should trigger the failing scheduler + ev1 = MagicMock() + ev1.unified_msg_origin = "event-1" + ev1.get_platform_id.return_value = "p1" + ev1.get_platform_name.return_value = "P1" + ev1.get_sender_name.return_value = None + ev1.get_sender_id.return_value = "u1" + ev1.get_message_outline.return_value = "m1" + + # Event 2 should trigger the working scheduler + ev2 = MagicMock() + ev2.unified_msg_origin = "event-2" + ev2.get_platform_id.return_value = "p2" + ev2.get_platform_name.return_value = "P2" + ev2.get_sender_name.return_value = None + ev2.get_sender_id.return_value = "u2" + ev2.get_message_outline.return_value = "m2" + + await event_queue.put(ev1) + await event_queue.put(ev2) + + task = asyncio.create_task(event_bus.dispatch()) + try: + await asyncio.wait_for(processed_second.wait(), timeout=2.0) + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + # Both schedulers should have been called + scheduler1.execute.assert_called_once_with(ev1) + scheduler2.execute.assert_called_once_with(ev2) + + @pytest.mark.asyncio + async def test_dispatch_continues_after_scheduler_exception( + self, event_queue, mock_config_manager + ): + """The dispatch loop must continue processing after a scheduler exception.""" + scheduler = MagicMock() + call_count = 0 + + async def execute_alternating(event): # noqa: ARG001 + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("First call fails") + + scheduler.execute.side_effect = execute_alternating + mock_config_manager.get_conf_info.return_value = { + "id": "same-conf", + "name": "Same", + } + mapping = {"same-conf": scheduler} + event_bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping=mapping, + astrbot_config_mgr=mock_config_manager, + ) + + ev = MagicMock() + ev.unified_msg_origin = "test:group:1" + ev.get_platform_id.return_value = "test" + ev.get_platform_name.return_value = "Test" + ev.get_sender_name.return_value = None + ev.get_sender_id.return_value = "u1" + ev.get_message_outline.return_value = "m" + + await event_queue.put(ev) + + task = asyncio.create_task(event_bus.dispatch()) + try: + await asyncio.wait_for(asyncio.sleep(0.2), timeout=1.0) + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + # scheduler should have been called once (event consumed, + # exception swallowed by asyncio.create_task) + assert call_count >= 1 + + +# ---- Edge-case inputs to dispatch ---- + + +class TestDispatchEdgeInputs: + """Dispatch handles unusual or missing event attributes.""" + + @pytest.mark.asyncio + async def test_dispatch_with_empty_origin( + self, event_queue, mock_config_manager, mock_scheduler + ): + """An event with empty unified_msg_origin is handled (falls back to '').""" + processed = asyncio.Event() + + mock_config_manager.get_conf_info.return_value = { + "id": "test-conf-id", + "name": "Test", + } + + async def execute_and_signal(event): # noqa: ARG001 + processed.set() + + mock_scheduler.execute.side_effect = execute_and_signal + + mapping = {"test-conf-id": mock_scheduler} + event_bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping=mapping, + astrbot_config_mgr=mock_config_manager, + ) + + ev = MagicMock() + ev.unified_msg_origin = "" # empty origin + ev.get_platform_id.return_value = "test" + ev.get_platform_name.return_value = "Test" + ev.get_sender_name.return_value = None + ev.get_sender_id.return_value = "u1" + ev.get_message_outline.return_value = "m" + + await event_queue.put(ev) + + task = asyncio.create_task(event_bus.dispatch()) + try: + await asyncio.wait_for(processed.wait(), timeout=1.0) + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + mock_config_manager.get_conf_info.assert_called_once_with("") + mock_scheduler.execute.assert_called_once() + + @pytest.mark.asyncio + async def test_dispatch_with_empty_config_info( + self, event_queue, mock_config_manager, mock_scheduler + ): + """get_conf_info returning empty dict still works (id falls back to '').""" + processed = asyncio.Event() + mock_config_manager.get_conf_info.return_value = {} + + async def execute_and_signal(event): # noqa: ARG001 + processed.set() + + mock_scheduler.execute.side_effect = execute_and_signal + + # The scheduler mapped to '' will be used (empty string from get("id", "")) + mapping = {"": mock_scheduler} + event_bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping=mapping, + astrbot_config_mgr=mock_config_manager, + ) + + ev = MagicMock() + ev.unified_msg_origin = "some:origin" + ev.get_platform_id.return_value = "test" + ev.get_platform_name.return_value = "Test" + ev.get_sender_name.return_value = None + ev.get_sender_id.return_value = "u1" + ev.get_message_outline.return_value = "m" + + await event_queue.put(ev) + + task = asyncio.create_task(event_bus.dispatch()) + try: + await asyncio.wait_for(processed.wait(), timeout=1.0) + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + mock_scheduler.execute.assert_called_once() + + @pytest.mark.asyncio + async def test_empty_pipeline_mapping_logs_error( + self, event_queue, mock_config_manager + ): + """An event is dropped with a logged error when no scheduler matches.""" + error_logged = asyncio.Event() + mock_config_manager.get_conf_info.return_value = { + "id": "orphan-id", + "name": "Orphan", + } + + event_bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping={}, + astrbot_config_mgr=mock_config_manager, + ) + + ev = MagicMock() + ev.unified_msg_origin = "test:private:1" + ev.get_platform_id.return_value = "test" + ev.get_platform_name.return_value = "Test" + ev.get_sender_name.return_value = "User" + ev.get_sender_id.return_value = "u1" + ev.get_message_outline.return_value = "m" + + await event_queue.put(ev) + + with patch("astrbot.core.event_bus.logger") as mock_logger: + mock_logger.error.side_effect = lambda *a, **kw: error_logged.set() + task = asyncio.create_task(event_bus.dispatch()) + try: + await asyncio.wait_for(error_logged.wait(), timeout=1.0) + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + assert "orphan-id" in mock_logger.error.call_args[0][0] + + +# ---- _print_event edge cases ---- + + +class TestPrintEventEdgeCases: + """_print_event handles unusual sender/platform names.""" + + def test_print_event_no_sender_name(self, event_bus_factory): + """When sender_name is None, the log omits the sender name section.""" + event_bus = event_bus_factory() + ev = MagicMock() + ev.get_platform_id.return_value = "test" + ev.get_platform_name.return_value = "TestPlatform" + ev.get_sender_name.return_value = None + ev.get_sender_id.return_value = "user-123" + ev.get_message_outline.return_value = "Hello" + + with patch("astrbot.core.event_bus.logger") as mock_logger: + event_bus._print_event(ev, "MyConfig") + + log_msg = mock_logger.info.call_args[0][0] + assert "MyConfig" in log_msg + assert "TestPlatform" in log_msg + assert "user-123" in log_msg + # Without sender name, there should be no '/' before user-123 + # (the format is [Config] [platform] sender_id: outline) + assert "Hello" in log_msg + + def test_print_event_sender_name_is_empty_string(self, event_bus_factory): + """An empty sender_name string is treated similarly to a missing name.""" + event_bus = event_bus_factory() + ev = MagicMock() + ev.get_platform_id.return_value = "test" + ev.get_platform_name.return_value = "TestPlatform" + ev.get_sender_name.return_value = "" + ev.get_sender_id.return_value = "user-123" + ev.get_message_outline.return_value = "Hello" + + with patch("astrbot.core.event_bus.logger") as mock_logger: + event_bus._print_event(ev, "MyConfig") + + mock_logger.info.assert_called_once() + + def test_print_event_with_special_characters(self, event_bus_factory): + """Special characters in event fields do not break logging.""" + event_bus = event_bus_factory() + ev = MagicMock() + ev.get_platform_id.return_value = "test@#$" + ev.get_platform_name.return_value = "Test[Platform]" + ev.get_sender_name.return_value = "User{}|" + ev.get_sender_id.return_value = "user:\n456" + ev.get_message_outline.return_value = "Hello\nWorld" + + with patch("astrbot.core.event_bus.logger") as mock_logger: + event_bus._print_event(ev, "Config[1]") + + mock_logger.info.assert_called_once() + + +# ---- EventBus construction edge cases ---- + + +class TestEventBusConstructionEdgeCases: + """EventBus handles edge values passed to __init__.""" + + def test_empty_pipeline_scheduler_mapping(self, event_queue, mock_config_manager): + """Empty pipeline_mapping is accepted at construction.""" + bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping={}, + astrbot_config_mgr=mock_config_manager, + ) + assert bus.pipeline_scheduler_mapping == {} + + def test_none_in_pipeline_mapping(self, event_queue, mock_config_manager): + """None values in the mapping are accepted (dealt with at dispatch time).""" + bus = EventBus( + event_queue=event_queue, + pipeline_scheduler_mapping={"conf-id": None}, + astrbot_config_mgr=mock_config_manager, + ) + assert bus.pipeline_scheduler_mapping["conf-id"] is None + + +# ---- Helper fixtures ---- + + +@pytest.fixture +def event_bus_factory(): + """Factory that creates EventBus with default test dependencies.""" + + def _create(event_queue=None, mapping=None, config_mgr=None): + import asyncio + + q = event_queue or asyncio.Queue() + m = mapping or {"test-id": MagicMock()} + cm = config_mgr or MagicMock() + return EventBus( + event_queue=q, + pipeline_scheduler_mapping=m, + astrbot_config_mgr=cm, + ) + + return _create diff --git a/tests/unit/test_file_extract.py b/tests/unit/test_file_extract.py new file mode 100644 index 000000000..065734d90 --- /dev/null +++ b/tests/unit/test_file_extract.py @@ -0,0 +1,267 @@ +"""Mock-based unit tests for astrbot.core.utils.file_extract. + +All external calls (httpx.AsyncClient and anyio.Path) are mocked so that +the tests never touch the network or the filesystem. +""" + +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from astrbot.core.utils.file_extract import extract_file_moonshotai + + +def _make_http_status_error(status_code: int = 400) -> httpx.HTTPStatusError: + """Return a realistic HTTPStatusError for use as a side_effect.""" + request = httpx.Request("POST", "https://api.moonshot.cn/v1/files") + response = httpx.Response(status_code, request=request) + return httpx.HTTPStatusError( + f"HTTP {status_code}", + request=request, + response=response, + ) + + +# --------------------------------------------------------------------------- +# mocks +# --------------------------------------------------------------------------- + + +def _mock_client() -> tuple[MagicMock, AsyncMock]: + """Create a pre-configured AsyncMock httpx client. + + Returns (mock_client, mock_httpx_cls) where: + - mock_httpx_cls.return_value = mock_client + - mock_client.__aenter__.return_value = mock_client + """ + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + return mock_client + + +def _mock_path(file_bytes: bytes = b"dummy file content") -> MagicMock: + """Create a pre-configured MagicMock anyio.Path. + + Returns mock_path where: + - mock_path.read_bytes is an AsyncMock returning *file_bytes* + """ + mock_path = MagicMock(spec=Path) # anyio.Path implements Path-like interface + mock_path.name = "test-file.pdf" # used for the upload filename + mock_path.read_bytes = AsyncMock(return_value=file_bytes) + return mock_path + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestExtractFileMoonshotaiSuccess: + + FAKE_CONTENT = b"%PDF-1.4 fake binary content..." + EXTRACTED_TEXT = "Extracted text from Moonshot AI API" + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_successful_extraction( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """Happy path: file is uploaded, file-id returned, content fetched.""" + mock_path = _mock_path(self.FAKE_CONTENT) + mock_path_cls.return_value = mock_path + + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + + upload_resp = MagicMock() + upload_resp.json.return_value = {"id": "file-mock-001"} + mock_client.post.return_value = upload_resp + + content_resp = MagicMock() + content_resp.text = self.EXTRACTED_TEXT + mock_client.get.return_value = content_resp + + result = await extract_file_moonshotai("/fake/path/report.pdf", "sk-test-key") + + assert result == self.EXTRACTED_TEXT + mock_path_cls.assert_called_once() + mock_path.read_bytes.assert_called_once() + mock_httpx_cls.assert_called_once() + mock_client.post.assert_called_once() + mock_client.get.assert_called_once_with("/files/file-mock-001/content") + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_client_created_with_authorization_header( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """The Bearer token should match the api_key passed.""" + mock_path_cls.return_value = _mock_path(b"data") + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + mock_client.post.return_value = MagicMock(json=lambda: {"id": "f1"}) + mock_client.get.return_value = MagicMock(text="ok") + + await extract_file_moonshotai("/f.pdf", "sk-secret-42") + + call_headers = mock_httpx_cls.call_args[1].get("headers", {}) + assert call_headers.get("Authorization") == "Bearer sk-secret-42" + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_passes_file_path_to_anyio_path( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """The file_path argument should be forwarded to anyio.Path().""" + mock_path_cls.return_value = _mock_path(b"data") + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + mock_client.post.return_value = MagicMock(json=lambda: {"id": "f1"}) + mock_client.get.return_value = MagicMock(text="ok") + + await extract_file_moonshotai("/my/custom/document.txt", "key") + + args, _ = mock_path_cls.call_args + assert str(args[0]) == "/my/custom/document.txt" + + +# --------------------------------------------------------------------------- +# Error paths +# --------------------------------------------------------------------------- + + +class TestExtractFileMoonshotaiErrors: + + FAKE_CONTENT = b"fake bytes" + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_missing_file_id_raises_value_error( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """Upload response without an 'id' key should raise ValueError.""" + mock_path_cls.return_value = _mock_path(self.FAKE_CONTENT) + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + + resp = MagicMock() + resp.json.return_value = {} # no "id" + mock_client.post.return_value = resp + + with pytest.raises(ValueError, match="valid file id"): + await extract_file_moonshotai("/f.pdf", "key") + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_null_file_id_raises_value_error( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """Upload response with 'id': None should raise ValueError.""" + mock_path_cls.return_value = _mock_path(self.FAKE_CONTENT) + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + + resp = MagicMock() + resp.json.return_value = {"id": None} + mock_client.post.return_value = resp + + with pytest.raises(ValueError, match="valid file id"): + await extract_file_moonshotai("/f.pdf", "key") + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_upload_http_error_propagates( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """Non-2xx status from upload POST should propagate as HTTPStatusError.""" + mock_path_cls.return_value = _mock_path(self.FAKE_CONTENT) + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + + resp = MagicMock() + resp.raise_for_status.side_effect = _make_http_status_error(400) + mock_client.post.return_value = resp + + with pytest.raises(httpx.HTTPStatusError): + await extract_file_moonshotai("/f.pdf", "key") + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_content_fetch_http_error_propagates( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """Non-2xx status from content GET should propagate as HTTPStatusError.""" + mock_path_cls.return_value = _mock_path(self.FAKE_CONTENT) + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + + upload_resp = MagicMock() + upload_resp.json.return_value = {"id": "file-001"} + mock_client.post.return_value = upload_resp + + content_resp = MagicMock() + content_resp.raise_for_status.side_effect = _make_http_status_error(500) + mock_client.get.return_value = content_resp + + with pytest.raises(httpx.HTTPStatusError): + await extract_file_moonshotai("/f.pdf", "key") + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_connection_error_during_upload_propagates( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """A network-level error on POST should propagate as ConnectError.""" + mock_path_cls.return_value = _mock_path(self.FAKE_CONTENT) + mock_client = _mock_client() + mock_httpx_cls.return_value = mock_client + + mock_client.post.side_effect = httpx.ConnectError("connection refused") + + with pytest.raises(httpx.ConnectError): + await extract_file_moonshotai("/f.pdf", "key") + + @pytest.mark.asyncio + @patch("astrbot.core.utils.file_extract.httpx.AsyncClient") + @patch("astrbot.core.utils.file_extract.anyio.Path") + async def test_file_read_error_prevents_http_call( + self, + mock_path_cls: MagicMock, + mock_httpx_cls: MagicMock, + ): + """If anyio.Path.read_bytes raises, the function should raise and + httpx.AsyncClient should never be constructed.""" + mock_path = _mock_path(b"ignored") + mock_path.read_bytes = AsyncMock(side_effect=OSError(2, "No such file")) + mock_path_cls.return_value = mock_path + + with pytest.raises(OSError): + await extract_file_moonshotai("/missing/file.txt", "key") + + mock_httpx_cls.assert_not_called() diff --git a/tests/unit/test_fixed_size_chunker.py b/tests/unit/test_fixed_size_chunker.py new file mode 100644 index 000000000..91b38ab35 --- /dev/null +++ b/tests/unit/test_fixed_size_chunker.py @@ -0,0 +1,314 @@ +""" +Unit tests for FixedSizeChunker. + +Covers construction, chunk method with various inputs (empty text, short text, +exact fit, overlap behavior, edge cases with chunk_size/overlap parameters, +and kwargs overriding instance defaults). +All tests isolate the chunker from any external dependencies. +""" + +import pytest + +from astrbot.core.knowledge_base.chunking.fixed_size import FixedSizeChunker + + +# --------------------------------------------------------------- +# Construction +# --------------------------------------------------------------- + + +class TestFixedSizeChunkerConstruction: + """Test construction of FixedSizeChunker.""" + + def test_default_construction(self): + """Test default parameters are set correctly.""" + chunker = FixedSizeChunker() + assert chunker.chunk_size == 512 + assert chunker.chunk_overlap == 50 + + def test_custom_construction(self): + """Test custom parameters are applied.""" + chunker = FixedSizeChunker(chunk_size=256, chunk_overlap=32) + assert chunker.chunk_size == 256 + assert chunker.chunk_overlap == 32 + + def test_zero_overlap_construction(self): + """Test that overlap of 0 is accepted.""" + chunker = FixedSizeChunker(chunk_size=100, chunk_overlap=0) + assert chunker.chunk_size == 100 + assert chunker.chunk_overlap == 0 + + def test_large_overlap_construction(self): + """Test that overlap >= chunk_size is accepted at construction.""" + chunker = FixedSizeChunker(chunk_size=10, chunk_overlap=10) + assert chunker.chunk_size == 10 + assert chunker.chunk_overlap == 10 + + +# --------------------------------------------------------------- +# chunk() - basic cases +# --------------------------------------------------------------- + + +class TestFixedSizeChunkerChunkBasic: + """Test basic chunk() behavior.""" + + @pytest.mark.asyncio + async def test_empty_text_returns_empty_list(self): + """Test that empty text returns an empty list.""" + chunker = FixedSizeChunker() + result = await chunker.chunk("") + assert result == [] + + @pytest.mark.asyncio + async def test_short_text_returns_single_chunk(self): + """Test that text shorter than chunk_size returns a single chunk.""" + chunker = FixedSizeChunker(chunk_size=100) + text = "Short text." + result = await chunker.chunk(text) + assert result == [text] + + @pytest.mark.asyncio + async def test_exact_fit_returns_single_chunk(self): + """Test that text exactly chunk_size returns a single chunk.""" + chunker = FixedSizeChunker(chunk_size=10) + text = "0123456789" + result = await chunker.chunk(text) + assert result == [text] + + @pytest.mark.asyncio + async def test_text_longer_than_chunk_size(self): + """Test that text longer than chunk_size is split into multiple chunks.""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=0) + text = "abcdefghij" + result = await chunker.chunk(text) + assert result == ["abcde", "fghij"] + + @pytest.mark.asyncio + async def test_whitespace_only_text(self): + """Test that whitespace-only text is handled without crashing.""" + chunker = FixedSizeChunker() + result = await chunker.chunk(" \n\n ") + assert isinstance(result, list) + + +# --------------------------------------------------------------- +# chunk() - overlap behavior +# --------------------------------------------------------------- + + +class TestFixedSizeChunkerOverlap: + """Test overlap behavior in chunk().""" + + @pytest.mark.asyncio + async def test_overlap_between_chunks(self): + """Test that chunks overlap correctly.""" + chunker = FixedSizeChunker(chunk_size=6, chunk_overlap=2) + text = "abcdefghij" + result = await chunker.chunk(text) + # start=0: "abcdef", start=4: "efghij" + assert result == ["abcdef", "efghij"] + + @pytest.mark.asyncio + async def test_overlap_equal_to_chunk_size(self): + """Test that when overlap >= chunk_size, the algorithm prevents an infinite loop.""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=5) + text = "abcdefghij" + result = await chunker.chunk(text) + # start=0: "abcde", since start >= end (5 >= 5), start becomes end + # no infinite loop; next: start=5: "fghij" + assert result == ["abcde", "fghij"] + + @pytest.mark.asyncio + async def test_overlap_greater_than_chunk_size(self): + """Test that when overlap > chunk_size, the algorithm prevents an infinite loop.""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=10) + text = "abcdefghij" + result = await chunker.chunk(text) + # start=0: "abcde", since start >= end (10 >= 5), start becomes end + # start=5: "fghij" + assert result == ["abcde", "fghij"] + + @pytest.mark.asyncio + async def test_no_overlap(self): + """Test that zero overlap produces non-overlapping chunks.""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=0) + text = "abcdefghij" + result = await chunker.chunk(text) + assert result == ["abcde", "fghij"] + + @pytest.mark.asyncio + async def test_partial_last_chunk_with_overlap(self): + """Test the last chunk is included even when shorter than chunk_size.""" + chunker = FixedSizeChunker(chunk_size=6, chunk_overlap=2) + text = "abcdefgh" + result = await chunker.chunk(text) + # start=0: "abcdef", start=4: "efgh" + assert result == ["abcdef", "efgh"] + + +# --------------------------------------------------------------- +# chunk() - kwargs override instance defaults +# --------------------------------------------------------------- + + +class TestFixedSizeChunkerKwargs: + """Test that kwargs override instance defaults in chunk().""" + + @pytest.mark.asyncio + async def test_chunk_size_kwarg_overrides_instance(self): + """Test that chunk_size in kwargs overrides the instance default.""" + chunker = FixedSizeChunker(chunk_size=500, chunk_overlap=0) + text = "abcdefghij" + result = await chunker.chunk(text, chunk_size=5) + assert result == ["abcde", "fghij"] + + @pytest.mark.asyncio + async def test_chunk_overlap_kwarg_overrides_instance(self): + """Test that chunk_overlap in kwargs overrides the instance default.""" + chunker = FixedSizeChunker(chunk_size=6, chunk_overlap=0) + text = "abcdefghij" + result = await chunker.chunk(text, chunk_overlap=2) + assert result == ["abcdef", "efghij"] + + @pytest.mark.asyncio + async def test_both_kwargs_override(self): + """Test that both kwargs override instance defaults.""" + chunker = FixedSizeChunker(chunk_size=100, chunk_overlap=10) + text = "abcdefghijklmnopqrstuvwxyz" + result = await chunker.chunk(text, chunk_size=10, chunk_overlap=3) + # start=0: "abcdefghij", start=7: "hijklmnopq", start=14: "opqrstuvwx", + # start=21: "vwxyz" + assert len(result) == 4 + + +# --------------------------------------------------------------- +# chunk() - edge cases +# --------------------------------------------------------------- + + +class TestFixedSizeChunkerEdgeCases: + """Test edge cases for chunk().""" + + @pytest.mark.asyncio + async def test_single_character_chunks(self): + """Test with chunk_size=1.""" + chunker = FixedSizeChunker(chunk_size=1, chunk_overlap=0) + text = "abc" + result = await chunker.chunk(text) + assert result == ["a", "b", "c"] + + @pytest.mark.asyncio + async def test_chunk_size_one_with_overlap(self): + """Test with chunk_size=1 and overlap=1 (should not infinite loop).""" + chunker = FixedSizeChunker(chunk_size=1, chunk_overlap=1) + text = "abc" + result = await chunker.chunk(text) + # start=0: "a", start>=end (1>=1) -> start=1, "b", ... + assert result == ["a", "b", "c"] + + @pytest.mark.asyncio + async def test_large_text(self): + """Test with a large chunk size that covers the entire text.""" + chunker = FixedSizeChunker(chunk_size=1000, chunk_overlap=50) + text = "A" * 900 + result = await chunker.chunk(text) + assert result == [text] + + @pytest.mark.asyncio + async def test_text_with_newlines(self): + """Test that newlines are handled as regular characters.""" + chunker = FixedSizeChunker(chunk_size=10, chunk_overlap=5) + text = "line1\nline2\nline3" + result = await chunker.chunk(text) + assert isinstance(result, list) + assert all(isinstance(c, str) for c in result) + # All characters should be accounted for + assert sum(len(c) for c in result) >= len(text) + + @pytest.mark.asyncio + async def test_unicode_text(self): + """Test that unicode characters are handled correctly.""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=2) + text = "你好世界abc" + result = await chunker.chunk(text) + # Chinese characters are 1 char each in Python + assert isinstance(result, list) + assert "".join(result) == text + + @pytest.mark.asyncio + async def test_exact_multiple_of_chunk_size(self): + """Test when text length is an exact multiple of chunk_size (no overlap).""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=0) + text = "abcde" * 3 # 15 chars + result = await chunker.chunk(text) + assert result == ["abcde", "fghij", "klmno"] + + @pytest.mark.asyncio + async def test_overlap_larger_than_chunk_minus_one(self): + """Test with overlap = chunk_size - 1 (maximum information overlap).""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=4) + text = "abcdefghij" + result = await chunker.chunk(text) + # start=0: "abcde", start=1: "bcdef", ..., start=5: "fghij" + assert result == ["abcde", "bcdef", "cdefg", "defgh", "efghi", "fghij"] + + @pytest.mark.asyncio + async def test_ten_thousand_characters(self): + """Test that chunker handles a moderately long string without issue.""" + chunker = FixedSizeChunker(chunk_size=100, chunk_overlap=10) + text = "x" * 10000 + result = await chunker.chunk(text) + # Expected length: ceil((10000 - 10) / (100 - 10)) = ceil(9990/90) = 111 + assert len(result) == 111 + assert all(len(c) == 100 or len(c) == 100 for c in result[:-1]) + assert len(result[-1]) <= 100 + + +# --------------------------------------------------------------- +# chunk() - verification of no content loss +# --------------------------------------------------------------- + + +class TestFixedSizeChunkerContentPreservation: + """Test that chunk() does not lose or reorder content.""" + + @pytest.mark.asyncio + async def test_no_content_loss_without_overlap(self): + """Test no content loss when overlap is 0.""" + chunker = FixedSizeChunker(chunk_size=10, chunk_overlap=0) + text = "0123456789" * 5 # 50 chars + result = await chunker.chunk(text) + combined = "".join(result) + assert combined == text + + @pytest.mark.asyncio + async def test_no_content_loss_with_overlap(self): + """Test that with overlap, original content is fully represented.""" + chunker = FixedSizeChunker(chunk_size=10, chunk_overlap=3) + text = "0123456789" * 5 + result = await chunker.chunk(text) + # With overlap, combined will be longer than original + combined = "".join(result) + assert len(combined) >= len(text) + for i, char in enumerate(text): + assert char in combined + + @pytest.mark.asyncio + async def test_chunks_are_substrings(self): + """Test that all chunks are valid substrings of the original text.""" + chunker = FixedSizeChunker(chunk_size=7, chunk_overlap=3) + text = "abcdefghijklmnop" + result = await chunker.chunk(text) + for chunk in result: + assert chunk in text + + @pytest.mark.asyncio + async def test_chunks_preserve_order(self): + """Test that chunks appear in the same order as the original text.""" + chunker = FixedSizeChunker(chunk_size=5, chunk_overlap=2) + text = "abcdefghij" + result = await chunker.chunk(text) + # Every chunk should be a substring and order should be consistent + positions = [text.index(chunk) for chunk in result] + assert positions == sorted(positions) diff --git a/tests/unit/test_kb_helper.py b/tests/unit/test_kb_helper.py new file mode 100644 index 000000000..607de7dc9 --- /dev/null +++ b/tests/unit/test_kb_helper.py @@ -0,0 +1,926 @@ +""" +Unit tests for KBHelper. + +Covers construction, initialize, get_ep, get_rp, _ensure_vec_db, terminate, +delete_vec_db, upload_document, list_documents, get_document, delete_document, +delete_chunk, refresh_kb, refresh_document, get_chunks_by_doc_id, +get_chunk_count_by_doc_id, _save_media, upload_from_url, +and _clean_and_rechunk_content. +All tests use mocks to isolate KBHelper from its dependencies. +""" + +import json +import sys +import types +from pathlib import Path +from unittest.mock import ANY, AsyncMock, MagicMock, PropertyMock, call, patch + +import pytest + + +@pytest.fixture +def stub_provider_manager_module(): + """Stub provider manager module to avoid circular imports in unit tests.""" + original_module = sys.modules.get("astrbot.core.provider.manager") + stub_module = types.ModuleType("astrbot.core.provider.manager") + + class ProviderManager: + ... + + setattr(stub_module, "ProviderManager", ProviderManager) + sys.modules["astrbot.core.provider.manager"] = stub_module + + try: + yield + finally: + if original_module is not None: + sys.modules["astrbot.core.provider.manager"] = original_module + else: + sys.modules.pop("astrbot.core.provider.manager", None) + + +@pytest.fixture +def mock_provider_manager(): + """Create a mock ProviderManager.""" + manager = MagicMock() + manager.get_provider_by_id = AsyncMock() + manager.acm = MagicMock() + manager.acm.default_conf = {} + return manager + + +@pytest.fixture +def mock_kb_db(): + """Create a mock KBSQLiteDatabase.""" + db = MagicMock() + db.get_db = MagicMock() + db.list_kbs = AsyncMock(return_value=[]) + db.get_kb_by_id = AsyncMock() + db.list_documents_by_kb = AsyncMock(return_value=[]) + db.get_document_by_id = AsyncMock() + db.delete_document_by_id = AsyncMock() + db.update_kb_stats = AsyncMock() + return db + + +@pytest.fixture +def mock_knowledge_base(): + """Create a mock KnowledgeBase instance using lazy import.""" + from astrbot.core.knowledge_base.models import KnowledgeBase + + kb = KnowledgeBase( + kb_name="test_kb", + description="Test knowledge base", + emoji="test", + embedding_provider_id="test-embedding-provider", + rerank_provider_id="test-rerank-provider", + chunk_size=512, + chunk_overlap=50, + top_k_dense=50, + top_k_sparse=50, + top_m_final=5, + ) + return kb + + +@pytest.fixture +def mock_embedding_provider(): + """Create a mock object that passes isinstance(EmbeddingProvider) checks.""" + from astrbot.core.provider.provider import EmbeddingProvider + + class FakeEmbeddingProvider(EmbeddingProvider): + def __init__(self) -> None: + pass + + async def get_embedding(self, text: str) -> list[float]: + return [0.1, 0.2, 0.3] + + async def get_embeddings(self, text: list[str]) -> list[list[float]]: + return [[0.1, 0.2, 0.3] for _ in text] + + def get_dim(self) -> int: + return 3 + + return FakeEmbeddingProvider() + + +@pytest.fixture +def mock_rerank_provider(): + """Create a mock object that passes isinstance(RerankProvider) checks.""" + from astrbot.core.provider.provider import RerankProvider + + class FakeRerankProvider(RerankProvider): + def __init__(self) -> None: + pass + + async def rerank_score(self, query: str, documents: list[str]) -> list[float]: + return [1.0] * len(documents) + + return FakeRerankProvider() + + +@pytest.fixture +def mock_chunker(): + """Create a mock chunker.""" + chunker = MagicMock() + chunker.chunk = AsyncMock(return_value=["chunk1", "chunk2"]) + return chunker + + +@pytest.fixture +def mock_vec_db(): + """Create a mock FaissVecDB.""" + vec_db = MagicMock() + vec_db.initialize = AsyncMock() + vec_db.close = AsyncMock() + vec_db.insert_batch = AsyncMock() + vec_db.delete = AsyncMock() + vec_db.count_documents = AsyncMock(return_value=5) + vec_db.document_storage = MagicMock() + vec_db.document_storage.get_documents = AsyncMock(return_value=[]) + return vec_db + + +@pytest.fixture +def helper_kwargs( + mock_kb_db, + mock_knowledge_base, + mock_provider_manager, + mock_chunker, +): + """Standard keyword arguments for constructing a KBHelper.""" + return { + "kb_db": mock_kb_db, + "kb": mock_knowledge_base, + "provider_manager": mock_provider_manager, + "kb_root_dir": "/tmp/test_kb_root", + "chunker": mock_chunker, + } + + +# --------------------------------------------------------------- +# Construction and initialization +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_construction_sets_attributes( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, + mock_knowledge_base, + mock_provider_manager, + mock_chunker, +): + """Test that KBHelper.__init__ sets all expected attributes.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper( + kb_db=mock_kb_db, + kb=mock_knowledge_base, + provider_manager=mock_provider_manager, + kb_root_dir="/tmp/test_kb_root", + chunker=mock_chunker, + ) + + assert helper.kb_db is mock_kb_db + assert helper.kb is mock_knowledge_base + assert helper.prov_mgr is mock_provider_manager + assert helper.chunker is mock_chunker + assert helper.init_error is None + assert helper.vec_db is None + assert isinstance(helper.kb_dir, Path) + assert "medias" in str(helper.kb_medias_dir) + assert "files" in str(helper.kb_files_dir) + + +@pytest.mark.asyncio +async def test_initialize_calls_ensure_vec_db( + stub_provider_manager_module, + helper_kwargs, +): + """Test that initialize() delegates to _ensure_vec_db.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + with ( + patch.object(KBHelper, "_ensure_vec_db", new_callable=AsyncMock) as mock_ensure, + ): + helper = KBHelper(**helper_kwargs) + await helper.initialize() + + mock_ensure.assert_awaited_once() + + +# --------------------------------------------------------------- +# get_ep +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_ep_raises_when_no_embedding_provider_id( + stub_provider_manager_module, + helper_kwargs, +): + """Test that get_ep raises ValueError when kb has no embedding_provider_id.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper(**helper_kwargs) + helper.kb.embedding_provider_id = None + + with pytest.raises(ValueError, match="未配置 Embedding Provider"): + await helper.get_ep() + + +@pytest.mark.asyncio +async def test_get_ep_raises_when_provider_not_found( + stub_provider_manager_module, + helper_kwargs, +): + """Test that get_ep raises ValueError when provider is not found by id.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper_kwargs["provider_manager"].get_provider_by_id.return_value = None + helper = KBHelper(**helper_kwargs) + + with pytest.raises(ValueError, match="无法找到"): + await helper.get_ep() + + +@pytest.mark.asyncio +async def test_get_ep_raises_when_not_embedding_provider( + stub_provider_manager_module, + helper_kwargs, + mock_provider_manager, +): + """Test that get_ep raises when the returned provider is not EmbeddingProvider.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = MagicMock() + helper = KBHelper(**helper_kwargs) + + with pytest.raises(ValueError, match="not an Embedding Provider"): + await helper.get_ep() + + +@pytest.mark.asyncio +async def test_get_ep_returns_embedding_provider( + stub_provider_manager_module, + helper_kwargs, + mock_embedding_provider, + mock_provider_manager, +): + """Test that get_ep returns the EmbeddingProvider successfully.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider + helper = KBHelper(**helper_kwargs) + + result = await helper.get_ep() + assert result is mock_embedding_provider + + +# --------------------------------------------------------------- +# get_rp +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_rp_returns_none_when_no_rerank_provider_id( + stub_provider_manager_module, + helper_kwargs, +): + """Test that get_rp returns None when kb has no rerank_provider_id.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper(**helper_kwargs) + helper.kb.rerank_provider_id = None + + result = await helper.get_rp() + assert result is None + + +@pytest.mark.asyncio +async def test_get_rp_returns_none_when_provider_not_found( + stub_provider_manager_module, + helper_kwargs, + mock_provider_manager, +): + """Test that get_rp returns None when the provider is not found (logs a warning).""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = None + helper = KBHelper(**helper_kwargs) + + result = await helper.get_rp() + assert result is None + + +@pytest.mark.asyncio +async def test_get_rp_raises_when_not_rerank_provider( + stub_provider_manager_module, + helper_kwargs, + mock_provider_manager, +): + """Test that get_rp raises when the returned provider is not RerankProvider.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = MagicMock() + helper = KBHelper(**helper_kwargs) + + with pytest.raises(ValueError, match="not a Rerank Provider"): + await helper.get_rp() + + +@pytest.mark.asyncio +async def test_get_rp_returns_rerank_provider( + stub_provider_manager_module, + helper_kwargs, + mock_rerank_provider, + mock_provider_manager, +): + """Test that get_rp returns the RerankProvider successfully.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_rerank_provider + helper = KBHelper(**helper_kwargs) + + result = await helper.get_rp() + assert result is mock_rerank_provider + + +# --------------------------------------------------------------- +# _ensure_vec_db +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_ensure_vec_db_creates_and_initializes( + stub_provider_manager_module, + helper_kwargs, + mock_embedding_provider, + mock_provider_manager, + mock_vec_db, +): + """Test that _ensure_vec_db creates FaissVecDB and initializes it.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider + helper = KBHelper(**helper_kwargs) + + with patch( + "astrbot.core.knowledge_base.kb_helper.FaissVecDB", + return_value=mock_vec_db, + ) as mock_faiss_cls: + result = await helper._ensure_vec_db() + + assert result is mock_vec_db + assert helper.vec_db is mock_vec_db + assert helper.init_error is None + mock_faiss_cls.assert_called_once() + mock_vec_db.initialize.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_ensure_vec_db_clears_stale_init_error( + stub_provider_manager_module, + helper_kwargs, + mock_embedding_provider, + mock_provider_manager, + mock_vec_db, +): + """Test that _ensure_vec_db clears a stale init_error on success.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider + helper = KBHelper(**helper_kwargs) + helper.init_error = "stale error" + + with patch( + "astrbot.core.knowledge_base.kb_helper.FaissVecDB", + return_value=mock_vec_db, + ): + await helper._ensure_vec_db() + assert helper.init_error is None + + +# --------------------------------------------------------------- +# terminate and delete_vec_db +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_terminate_closes_vec_db( + stub_provider_manager_module, + helper_kwargs, + mock_vec_db, +): + """Test that terminate() closes the vec_db.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper(**helper_kwargs) + helper.vec_db = mock_vec_db + + await helper.terminate() + mock_vec_db.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_terminate_handles_missing_vec_db( + stub_provider_manager_module, + helper_kwargs, +): + """Test that terminate() does not crash when vec_db is None.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper(**helper_kwargs) + helper.vec_db = None + + # Should not raise + await helper.terminate() + + +# --------------------------------------------------------------- +# upload_document +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_upload_document_with_pre_chunked_text( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, + mock_vec_db, + mock_embedding_provider, + mock_provider_manager, +): + """Test that upload_document works with pre-chunked text (skips parsing).""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider + helper = KBHelper(**helper_kwargs) + + # Mock the async session for document metadata persistence + mock_session = MagicMock() + mock_session.add = MagicMock() + mock_session.begin = MagicMock() + mock_session.begin.return_value.__aenter__ = AsyncMock() + mock_session.begin.return_value.__aexit__ = AsyncMock() + mock_session.refresh = AsyncMock() + mock_db_ctx = MagicMock() + mock_db_ctx.__aenter__ = AsyncMock(return_value=mock_session) + mock_db_ctx.__aexit__ = AsyncMock() + mock_kb_db.get_db.return_value = mock_db_ctx + + with patch.object(helper, "_get_vec_db", return_value=mock_vec_db): + doc = await helper.upload_document( + file_name="test.txt", + file_content=None, + file_type="txt", + pre_chunked_text=["chunk a", "chunk b"], + ) + + assert doc.doc_name == "test.txt" + assert doc.file_type == "txt" + # Metadata should have been added + mock_session.add.assert_called() + mock_vec_db.insert_batch.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_upload_document_parses_when_no_pre_chunked( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, + mock_vec_db, + mock_embedding_provider, + mock_provider_manager, + mock_chunker, +): + """Test that upload_document parses file content when no pre-chunked text.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider + mock_chunker.chunk = AsyncMock(return_value=["parsed chunk"]) + helper = KBHelper(**helper_kwargs) + + mock_session = MagicMock() + mock_session.add = MagicMock() + mock_session.begin = MagicMock() + mock_session.begin.return_value.__aenter__ = AsyncMock() + mock_session.begin.return_value.__aexit__ = AsyncMock() + mock_session.refresh = AsyncMock() + mock_db_ctx = MagicMock() + mock_db_ctx.__aenter__ = AsyncMock(return_value=mock_session) + mock_db_ctx.__aexit__ = AsyncMock() + mock_kb_db.get_db.return_value = mock_db_ctx + + mock_parser = MagicMock() + mock_parse_result = MagicMock() + mock_parse_result.text = "some parsed text" + mock_parse_result.media = [] + mock_parser.parse = AsyncMock(return_value=mock_parse_result) + + with ( + patch.object(helper, "_get_vec_db", return_value=mock_vec_db), + patch( + "astrbot.core.knowledge_base.kb_helper.select_parser", + return_value=mock_parser, + ), + ): + doc = await helper.upload_document( + file_name="test.pdf", + file_content=b"%PDF-1.4 fake content", + file_type="pdf", + ) + + assert doc.doc_name == "test.pdf" + mock_parser.parse.assert_awaited_once() + mock_chunker.chunk.assert_awaited_once() + mock_vec_db.insert_batch.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_upload_document_raises_when_no_content_and_no_pre_chunked( + stub_provider_manager_module, + helper_kwargs, + mock_vec_db, + mock_embedding_provider, + mock_provider_manager, +): + """Test that upload_document raises when file_content is None and no pre_chunked.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider + helper = KBHelper(**helper_kwargs) + + with ( + patch.object(helper, "_get_vec_db", return_value=mock_vec_db), + pytest.raises(ValueError, match="file_content 不能为空"), + ): + await helper.upload_document( + file_name="test.pdf", + file_content=None, + file_type="pdf", + ) + + +# --------------------------------------------------------------- +# list_documents / get_document +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_list_documents_delegates_to_db( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, +): + """Test that list_documents delegates to kb_db.list_documents_by_kb.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.list_documents_by_kb.return_value = ["doc1", "doc2"] + helper = KBHelper(**helper_kwargs) + + result = await helper.list_documents(offset=0, limit=50) + + mock_kb_db.list_documents_by_kb.assert_awaited_once_with( + helper.kb.kb_id, + 0, + 50, + ) + assert result == ["doc1", "doc2"] + + +@pytest.mark.asyncio +async def test_get_document_delegates_to_db( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, +): + """Test that get_document delegates to kb_db.get_document_by_id.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.get_document_by_id.return_value = "fake_doc" + helper = KBHelper(**helper_kwargs) + + result = await helper.get_document("doc-123") + + mock_kb_db.get_document_by_id.assert_awaited_once_with("doc-123") + assert result == "fake_doc" + + +# --------------------------------------------------------------- +# delete_document / delete_chunk +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_delete_document_delegates_to_db_and_vec_db( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, + mock_vec_db, +): + """Test that delete_document cleans up via kb_db and vec_db.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.get_kb_by_id = AsyncMock(return_value=helper_kwargs["kb"]) + helper = KBHelper(**helper_kwargs) + helper.vec_db = mock_vec_db + + await helper.delete_document("doc-123") + + mock_kb_db.delete_document_by_id.assert_awaited_once_with( + doc_id="doc-123", + vec_db=mock_vec_db, + ) + mock_kb_db.update_kb_stats.assert_awaited() + + +@pytest.mark.asyncio +async def test_delete_chunk_delegates_to_vec_db( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, + mock_vec_db, +): + """Test that delete_chunk removes a chunk from vec_db and refreshes stats.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.get_kb_by_id = AsyncMock(return_value=helper_kwargs["kb"]) + helper = KBHelper(**helper_kwargs) + helper.vec_db = mock_vec_db + + await helper.delete_chunk("chunk-1", "doc-123") + + mock_vec_db.delete.assert_awaited_once_with("chunk-1") + + +# --------------------------------------------------------------- +# refresh_kb / refresh_document +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_refresh_kb_updates_self_kb( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, + mock_knowledge_base, +): + """Test that refresh_kb fetches the latest KB from the database.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + fresh_kb = MagicMock() + fresh_kb.kb_id = mock_knowledge_base.kb_id + mock_kb_db.get_kb_by_id = AsyncMock(return_value=fresh_kb) + + helper = KBHelper(**helper_kwargs) + + await helper.refresh_kb() + + mock_kb_db.get_kb_by_id.assert_awaited_once_with(mock_knowledge_base.kb_id) + assert helper.kb is fresh_kb + + +@pytest.mark.asyncio +async def test_refresh_document_raises_when_not_found( + stub_provider_manager_module, + helper_kwargs, + mock_kb_db, +): + """Test that refresh_document raises ValueError when doc does not exist.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.get_document_by_id = AsyncMock(return_value=None) + helper = KBHelper(**helper_kwargs) + + with pytest.raises(ValueError, match="无法找到"): + await helper.refresh_document("doc-123") + + +# --------------------------------------------------------------- +# get_chunks_by_doc_id +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_chunks_by_doc_id_returns_formatted_chunks( + stub_provider_manager_module, + helper_kwargs, + mock_vec_db, +): + """Test that get_chunks_by_doc_id returns properly formatted chunks.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_vec_db.document_storage.get_documents.return_value = [ + { + "doc_id": "chunk-1", + "metadata": json.dumps( + {"kb_doc_id": "doc-123", "kb_id": "kb-1", "chunk_index": 0}, + ), + "text": "chunk content", + }, + ] + helper = KBHelper(**helper_kwargs) + helper.vec_db = mock_vec_db + + result = await helper.get_chunks_by_doc_id("doc-123") + + assert len(result) == 1 + chunk = result[0] + assert chunk["chunk_id"] == "chunk-1" + assert chunk["doc_id"] == "doc-123" + assert chunk["content"] == "chunk content" + assert chunk["char_count"] == 13 + + +# --------------------------------------------------------------- +# get_chunk_count_by_doc_id +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_chunk_count_by_doc_id( + stub_provider_manager_module, + helper_kwargs, + mock_vec_db, +): + """Test that get_chunk_count_by_doc_id returns the count from vec_db.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_vec_db.count_documents.return_value = 42 + helper = KBHelper(**helper_kwargs) + helper.vec_db = mock_vec_db + + count = await helper.get_chunk_count_by_doc_id("doc-123") + + assert count == 42 + mock_vec_db.count_documents.assert_awaited_once_with( + metadata_filter={"kb_doc_id": "doc-123"}, + ) + + +# --------------------------------------------------------------- +# _save_media +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_save_media_creates_media_record( + stub_provider_manager_module, + helper_kwargs, +): + """Test that _save_media saves the media file and returns a KBMedia record.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper(**helper_kwargs) + + with patch("astrbot.core.knowledge_base.kb_helper.aiofiles.open") as mock_aiofiles: + mock_file = AsyncMock() + mock_aiofiles.return_value.__aenter__ = AsyncMock(return_value=mock_file) + mock_aiofiles.return_value.__aexit__ = AsyncMock() + + media = await helper._save_media( + doc_id="doc-1", + media_type="image", + file_name="photo.png", + content=b"fake-image-bytes", + mime_type="image/png", + ) + + assert media.media_type == "image" + assert media.file_name == "photo.png" + assert media.mime_type == "image/png" + assert media.doc_id == "doc-1" + mock_file.write.assert_awaited_once_with(b"fake-image-bytes") + + +# --------------------------------------------------------------- +# upload_from_url +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_upload_from_url_raises_when_no_tavily_key( + stub_provider_manager_module, + helper_kwargs, + mock_provider_manager, +): + """Test that upload_from_url raises when Tavily API key is missing.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.acm.default_conf = {"provider_settings": {}} + helper = KBHelper(**helper_kwargs) + + with pytest.raises(ValueError, match="Tavily API key"): + await helper.upload_from_url("http://example.com") + + +@pytest.mark.asyncio +async def test_upload_from_url_delegates_to_upload_document( + stub_provider_manager_module, + helper_kwargs, + mock_provider_manager, +): + """Test that upload_from_url extracts text and delegates to upload_document.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.acm.default_conf = { + "provider_settings": { + "websearch_tavily_key": ["fake-key"], + }, + } + helper = KBHelper(**helper_kwargs) + + with ( + patch( + "astrbot.core.knowledge_base.kb_helper.extract_text_from_url", + new_callable=AsyncMock, + return_value="extracted content", + ), + patch.object( + helper, + "upload_document", + new_callable=AsyncMock, + return_value="fake_doc", + ), + ): + result = await helper.upload_from_url("http://example.com") + + assert result == "fake_doc" + + +# --------------------------------------------------------------- +# _clean_and_rechunk_content +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_clean_and_rechunk_skips_when_not_enabled( + stub_provider_manager_module, + helper_kwargs, + mock_chunker, +): + """Test that _clean_and_rechunk_content uses chunker directly when cleaning disabled.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_chunker.chunk = AsyncMock(return_value=["chunk a", "chunk b"]) + helper = KBHelper(**helper_kwargs) + + result = await helper._clean_and_rechunk_content( + content="some text", + url="http://example.com", + enable_cleaning=False, + ) + + assert result == ["chunk a", "chunk b"] + mock_chunker.chunk.assert_awaited_once_with("some text") + + +@pytest.mark.asyncio +async def test_clean_and_rechunk_skips_when_no_provider_id( + stub_provider_manager_module, + helper_kwargs, + mock_chunker, +): + """Test that _clean_and_rechunk_content uses chunker when cleaning_provider_id is None.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_chunker.chunk = AsyncMock(return_value=["default chunk"]) + helper = KBHelper(**helper_kwargs) + + result = await helper._clean_and_rechunk_content( + content="some text", + url="http://example.com", + enable_cleaning=True, + cleaning_provider_id=None, + ) + + assert result == ["default chunk"] + + +@pytest.mark.asyncio +async def test_clean_and_rechunk_falls_back_on_error( + stub_provider_manager_module, + helper_kwargs, + mock_provider_manager, + mock_chunker, +): + """Test that _clean_and_rechunk_content falls back to chunker when LLM call fails.""" + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_provider_manager.get_provider_by_id = AsyncMock( + side_effect=Exception("provider error"), + ) + mock_chunker.chunk = AsyncMock(return_value=["fallback chunk"]) + helper = KBHelper(**helper_kwargs) + + result = await helper._clean_and_rechunk_content( + content="some text", + url="http://example.com", + enable_cleaning=True, + cleaning_provider_id="llm-1", + ) + + assert result == ["fallback chunk"] diff --git a/tests/unit/test_kb_mgr.py b/tests/unit/test_kb_mgr.py new file mode 100644 index 000000000..f6ed94826 --- /dev/null +++ b/tests/unit/test_kb_mgr.py @@ -0,0 +1,657 @@ +""" +Unit tests for KnowledgeBaseManager. + +Covers construction, initialize, create_kb, get_kb, get_kb_by_name, +delete_kb, list_kbs, update_kb, retrieve, terminate, and upload_from_url. +All tests use mocks to isolate the manager from its dependencies. +""" + +import sys +import types +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch + +import pytest + + +@pytest.fixture +def stub_provider_manager_module(): + """Stub provider manager module to avoid circular imports in unit tests.""" + original_module = sys.modules.get("astrbot.core.provider.manager") + stub_module = types.ModuleType("astrbot.core.provider.manager") + + class ProviderManager: + ... + + setattr(stub_module, "ProviderManager", ProviderManager) + sys.modules["astrbot.core.provider.manager"] = stub_module + + try: + yield + finally: + if original_module is not None: + sys.modules["astrbot.core.provider.manager"] = original_module + else: + sys.modules.pop("astrbot.core.provider.manager", None) + + +@pytest.fixture +def mock_provider_manager(): + """Create a mock ProviderManager.""" + manager = MagicMock() + manager.get_provider_by_id = AsyncMock() + manager.acm = MagicMock() + manager.acm.default_conf = {} + return manager + + +@pytest.fixture +def mock_kb_db(): + """Create a mock KBSQLiteDatabase.""" + db = MagicMock() + db.get_db = MagicMock() + db.list_kbs = AsyncMock(return_value=[]) + db.get_kb_by_id = AsyncMock() + return db + + +@pytest.fixture +def mock_session(): + """Create a mock async session with transaction helpers.""" + session = MagicMock() + session.add = MagicMock() + session.flush = AsyncMock() + session.commit = AsyncMock() + session.refresh = AsyncMock() + session.begin = MagicMock() + session.begin.return_value.__aenter__ = AsyncMock() + session.begin.return_value.__aexit__ = AsyncMock() + session.delete = AsyncMock() + return session + + +@pytest.fixture +def mock_db_context(mock_session): + """Create a mock async context manager for get_db().""" + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=mock_session) + ctx.__aexit__ = AsyncMock() + return ctx + + +@pytest.fixture +def mock_knowledge_base(): + """Create a mock KnowledgeBase instance using lazy import.""" + from astrbot.core.knowledge_base.models import KnowledgeBase + + kb = KnowledgeBase( + kb_name="test_kb", + description="Test knowledge base", + emoji="test", + embedding_provider_id="test-embedding-provider", + rerank_provider_id=None, + chunk_size=512, + chunk_overlap=50, + top_k_dense=50, + top_k_sparse=50, + top_m_final=5, + ) + return kb + + +# --------------------------------------------------------------- +# Construction & initialization +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_manager_construction(stub_provider_manager_module, mock_provider_manager): + """Test that KnowledgeBaseManager can be constructed with a provider manager.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.provider_manager = mock_provider_manager + mgr.kb_insts = {} + mgr._session_deleted_callback_registered = False + + assert mgr.provider_manager is mock_provider_manager + assert mgr.kb_insts == {} + assert mgr._session_deleted_callback_registered is False + + +@pytest.mark.asyncio +async def test_manager_initialize_creates_db_and_loads_kbs( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, +): + """Test that initialize() creates the database and loads existing KBs.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.provider_manager = mock_provider_manager + mgr.kb_insts = {} + + with ( + patch( + "astrbot.core.knowledge_base.kb_mgr.KBSQLiteDatabase", + return_value=mock_kb_db, + ) as mock_db_cls, + patch( + "astrbot.core.knowledge_base.kb_mgr.RetrievalManager", + ) as mock_retrieval_cls, + patch( + "astrbot.core.knowledge_base.kb_mgr.SparseRetriever", + ), + patch( + "astrbot.core.knowledge_base.kb_mgr.RankFusion", + ), + ): + mock_retrieval = MagicMock() + mock_retrieval_cls.return_value = mock_retrieval + + await mgr.initialize() + + mock_db_cls.assert_called_once() + mock_kb_db.initialize.assert_awaited_once() + mock_kb_db.migrate_to_v1.assert_awaited_once() + mock_kb_db.list_kbs.assert_awaited_once() + assert mgr.retrieval_manager is mock_retrieval + + +@pytest.mark.asyncio +async def test_initialize_handles_import_error_gracefully( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that initialize() catches ImportError without crashing.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.provider_manager = mock_provider_manager + mgr.kb_insts = {} + + with patch( + "astrbot.core.knowledge_base.kb_mgr.KBSQLiteDatabase", + side_effect=ImportError("missing dependency"), + ): + await mgr.initialize() + # Should not raise — the error is logged + assert not hasattr(mgr, "retrieval_manager") or mgr.retrieval_manager is None + + +# --------------------------------------------------------------- +# create_kb +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_kb_raises_when_embedding_provider_id_is_none( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, +): + """Test that create_kb raises ValueError when embedding_provider_id is None.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.provider_manager = mock_provider_manager + mgr.kb_db = mock_kb_db + mgr.kb_insts = {} + + with pytest.raises(ValueError, match="embedding_provider_id"): + await mgr.create_kb(kb_name="my_kb") + + +@pytest.mark.asyncio +async def test_create_kb_success( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, + mock_db_context, + mock_session, +): + """Test that create_kb creates a new KB, persists it, and returns a KBHelper.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.get_db.return_value = mock_db_context + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.provider_manager = mock_provider_manager + mgr.kb_db = mock_kb_db + mgr.kb_insts = {} + + with patch.object(KBHelper, "initialize", new_callable=AsyncMock) as mock_init: + mock_init.return_value = None + + result = await mgr.create_kb( + kb_name="my_kb", + description="desc", + emoji="doc", + embedding_provider_id="ep-1", + ) + + assert result is not None + assert result.kb.kb_name == "my_kb" + assert result.kb.embedding_provider_id == "ep-1" + mock_session.add.assert_called_once() + mock_session.flush.assert_awaited_once() + mock_init.assert_awaited_once() + mock_session.commit.assert_awaited_once() + assert result.kb.kb_id in mgr.kb_insts + + +@pytest.mark.asyncio +async def test_create_kb_duplicate_name_raises( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, + mock_db_context, + mock_session, +): + """Test that create_kb raises ValueError on duplicate kb_name.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mock_kb_db.get_db.return_value = mock_db_context + # Simulate an IntegrityError-like message + mock_session.flush = AsyncMock( + side_effect=Exception("UNIQUE constraint failed: kb_name"), + ) + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.provider_manager = mock_provider_manager + mgr.kb_db = mock_kb_db + mgr.kb_insts = {} + + with pytest.raises(ValueError, match="已存在"): + await mgr.create_kb( + kb_name="my_kb", + embedding_provider_id="ep-1", + ) + + +# --------------------------------------------------------------- +# get_kb / get_kb_by_name +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_kb_returns_none_for_unknown_id( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that get_kb returns None when kb_id is not found.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {} + + result = await mgr.get_kb("nonexistent") + assert result is None + + +@pytest.mark.asyncio +async def test_get_kb_returns_helper_for_known_id( + stub_provider_manager_module, + mock_provider_manager, + mock_knowledge_base, +): + """Test that get_kb returns the correct KBHelper for a known kb_id.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + result = await mgr.get_kb(mock_knowledge_base.kb_id) + assert result is helper + + +@pytest.mark.asyncio +async def test_get_kb_by_name_returns_helper( + stub_provider_manager_module, + mock_provider_manager, + mock_knowledge_base, +): + """Test that get_kb_by_name returns the correct helper by matching kb_name.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + result = await mgr.get_kb_by_name("test_kb") + assert result is helper + + +@pytest.mark.asyncio +async def test_get_kb_by_name_returns_none_for_missing( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that get_kb_by_name returns None when no match is found.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {} + + result = await mgr.get_kb_by_name("nonexistent") + assert result is None + + +# --------------------------------------------------------------- +# delete_kb +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_delete_kb_returns_false_for_unknown( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that delete_kb returns False when kb_id is not found.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {} + + result = await mgr.delete_kb("nonexistent") + assert result is False + + +@pytest.mark.asyncio +async def test_delete_kb_removes_helper( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, + mock_db_context, + mock_session, + mock_knowledge_base, +): + """Test that delete_kb removes the KBHelper, deletes vec_db, and removes from DB.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + mock_kb_db.get_db.return_value = mock_db_context + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + helper.vec_db = MagicMock() + helper.delete_vec_db = AsyncMock() + helper.terminate = AsyncMock() + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_db = mock_kb_db + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + result = await mgr.delete_kb(mock_knowledge_base.kb_id) + + assert result is True + helper.delete_vec_db.assert_awaited_once() + mock_session.delete.assert_awaited_once_with(mock_knowledge_base) + mock_session.commit.assert_awaited_once() + assert mock_knowledge_base.kb_id not in mgr.kb_insts + + +# --------------------------------------------------------------- +# list_kbs +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_list_kbs_returns_empty_when_no_instances( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that list_kbs returns an empty list when no KBs exist.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {} + + result = await mgr.list_kbs() + assert result == [] + + +@pytest.mark.asyncio +async def test_list_kbs_returns_all_knowledge_bases( + stub_provider_manager_module, + mock_provider_manager, + mock_knowledge_base, +): + """Test that list_kbs returns KnowledgeBase objects for all instances.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + result = await mgr.list_kbs() + assert len(result) == 1 + assert result[0] is mock_knowledge_base + + +# --------------------------------------------------------------- +# terminate +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_terminate_closes_all_helpers_and_db( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, + mock_knowledge_base, +): + """Test that terminate() terminates all helpers and closes the database.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + helper.terminate = AsyncMock() + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_db = mock_kb_db + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + mock_kb_db.close = AsyncMock() + + await mgr.terminate() + + helper.terminate.assert_awaited_once() + mock_kb_db.close.assert_awaited_once() + assert mgr.kb_insts == {} + + +@pytest.mark.asyncio +async def test_terminate_handles_helper_failure_gracefully( + stub_provider_manager_module, + mock_provider_manager, + mock_kb_db, + mock_knowledge_base, +): + """Test that terminate() continues even if one helper raises.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + helper.terminate = AsyncMock(side_effect=Exception("close failed")) + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_db = mock_kb_db + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + mock_kb_db.close = AsyncMock() + + # Should not raise + await mgr.terminate() + helper.terminate.assert_awaited_once() + mock_kb_db.close.assert_awaited_once() + + +# --------------------------------------------------------------- +# retrieve +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_retrieve_returns_empty_when_no_kb_found( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that retrieve returns an empty dict when no KBs match the given names.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {} + + result = await mgr.retrieve("query", kb_names=["nonexistent"]) + assert result == {} + + +@pytest.mark.asyncio +async def test_retrieve_raises_when_all_kbs_unavailable( + stub_provider_manager_module, + mock_provider_manager, + mock_knowledge_base, +): + """Test that retrieve raises ValueError when all matching KBs have init_error.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + helper.init_error = "provider unavailable" + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + with pytest.raises(ValueError, match="所有请求的知识库均不可用"): + await mgr.retrieve("query", kb_names=["test_kb"]) + + +@pytest.mark.asyncio +async def test_retrieve_returns_formatted_results( + stub_provider_manager_module, + mock_provider_manager, + mock_knowledge_base, +): + """Test that retrieve returns properly formatted results when KBs are available.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + helper.init_error = None + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + mgr.kb_db = MagicMock() + + retrieval_manager = MagicMock() + retrieval_manager.retrieve = AsyncMock() + + from collections import namedtuple + + FakeResult = namedtuple( + "FakeResult", + [ + "chunk_id", + "doc_id", + "kb_id", + "kb_name", + "doc_name", + "chunk_index", + "content", + "score", + "metadata", + ], + ) + fake_result = FakeResult( + chunk_id="c1", + doc_id="d1", + kb_id=mock_knowledge_base.kb_id, + kb_name="test_kb", + doc_name="doc1.pdf", + chunk_index=0, + content="some content", + score=0.95, + metadata={"chunk_index": 0, "char_count": 12}, + ) + retrieval_manager.retrieve.return_value = [fake_result] + mgr.retrieval_manager = retrieval_manager + + result = await mgr.retrieve("query", kb_names=["test_kb"]) + + assert result is not None + assert "context_text" in result + assert "results" in result + assert len(result["results"]) == 1 + assert result["results"][0]["kb_name"] == "test_kb" + assert result["results"][0]["content"] == "some content" + retrieval_manager.retrieve.assert_awaited_once() + + +# --------------------------------------------------------------- +# upload_from_url +# --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_upload_from_url_raises_when_kb_not_found( + stub_provider_manager_module, + mock_provider_manager, +): + """Test that upload_from_url raises ValueError when kb_id is not found.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {} + + with pytest.raises(ValueError, match="not found"): + await mgr.upload_from_url("nonexistent", "http://example.com") + + +@pytest.mark.asyncio +async def test_upload_from_url_delegates_to_helper( + stub_provider_manager_module, + mock_provider_manager, + mock_knowledge_base, +): + """Test that upload_from_url delegates to the correct KBHelper.""" + from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + from astrbot.core.knowledge_base.kb_helper import KBHelper + + helper = KBHelper.__new__(KBHelper) + helper.kb = mock_knowledge_base + helper.upload_from_url = AsyncMock(return_value="fake_doc") + + mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager) + mgr.kb_insts = {mock_knowledge_base.kb_id: helper} + + result = await mgr.upload_from_url( + mock_knowledge_base.kb_id, + "http://example.com", + ) + + assert result == "fake_doc" + helper.upload_from_url.assert_awaited_once_with( + url="http://example.com", + chunk_size=512, + chunk_overlap=50, + batch_size=32, + tasks_limit=3, + max_retries=3, + progress_callback=None, + ) diff --git a/tests/unit/test_long_term_memory.py b/tests/unit/test_long_term_memory.py new file mode 100644 index 000000000..c1a7d06ac --- /dev/null +++ b/tests/unit/test_long_term_memory.py @@ -0,0 +1,283 @@ +"""Mock-based unit tests for LongTermMemory.""" + +from __future__ import annotations + +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest + +from astrbot.builtin_stars.astrbot.long_term_memory import LongTermMemory +from astrbot.core.platform.message_type import MessageType + + +@pytest.fixture +def mock_acm(): + return MagicMock() + + +@pytest.fixture +def mock_context(): + return MagicMock() + + +@pytest.fixture +def ltm(mock_acm, mock_context): + return LongTermMemory(mock_acm, mock_context) + + +@pytest.fixture +def mock_event(): + event = MagicMock() + event.unified_msg_origin = "qq:group:123456" + event.get_message_type.return_value = MessageType.GROUP_MESSAGE + event.is_at_or_wake_command = False + event.get_group_id.return_value = "123456" + event.get_messages.return_value = [] + event.message_obj.sender.nickname = "TestUser" + return event + + +class TestLongTermMemoryConstruction: + """Construction and initial state.""" + + def test_init_stores_deps(self, ltm, mock_acm, mock_context): + assert ltm.acm is mock_acm + assert ltm.context is mock_context + + def test_init_session_chats_is_defaultdict(self, ltm): + assert ltm.session_chats["any_key"] == [] + + @patch.object(LongTermMemory, "cfg", return_value={"max_cnt": 100, "image_caption": False, "enable_active_reply": False}) + def test_remove_session_returns_count(self, mock_cfg, ltm, mock_event): + ltm.session_chats["qq:group:123456"] = ["msg1", "msg2"] + cnt = 0 + import asyncio + cnt = asyncio.run(ltm.remove_session(mock_event)) + assert cnt == 2 + assert "qq:group:123456" not in ltm.session_chats + + +class TestCfg: + """Configuration extraction from context.""" + + def test_cfg_reads_from_context(self, ltm, mock_context, mock_event): + fake_ctx_cfg = { + "provider_ltm_settings": { + "group_message_max_cnt": "500", + "image_caption": True, + "image_caption_provider_id": "prov-1", + "active_reply": { + "enable": True, + "method": "possibility_reply", + "possibility_reply": 0.5, + "whitelist": [], + }, + }, + "provider_settings": { + "image_caption_prompt": "Describe", + }, + } + mock_context.get_config.return_value = fake_ctx_cfg + result = ltm.cfg(mock_event) + assert result["max_cnt"] == 500 + assert result["image_caption"] is True + assert result["enable_active_reply"] is True + assert result["image_caption_prompt"] == "Describe" + + def test_cfg_handles_missing_max_cnt(self, ltm, mock_context, mock_event): + mock_context.get_config.return_value = { + "provider_ltm_settings": {}, + "provider_settings": {"image_caption_prompt": "x"}, + } + result = ltm.cfg(mock_event) + assert result["max_cnt"] == 300 # default fallback + assert result["image_caption"] is False + + +class TestGetImageCaption: + """Image caption fetching.""" + + async def test_get_image_caption_uses_using_provider(self, ltm, mock_context): + mock_provider = AsyncMock() + mock_provider.text_chat = AsyncMock(return_value=MagicMock(completion_text="a cat")) + mock_context.get_using_provider.return_value = mock_provider + mock_context.get_provider_by_id.return_value = None + + caption = await ltm.get_image_caption("http://img.jpg", "", "Describe") + assert caption == "a cat" + mock_provider.text_chat.assert_awaited_once() + + async def test_get_image_caption_uses_provider_by_id(self, ltm, mock_context): + mock_provider = AsyncMock() + mock_provider.text_chat = AsyncMock(return_value=MagicMock(completion_text="a dog")) + mock_context.get_provider_by_id.return_value = mock_provider + + caption = await ltm.get_image_caption("http://img.jpg", "custom-provider", "Describe") + assert caption == "a dog" + + async def test_get_image_caption_raises_on_missing_provider(self, ltm, mock_context): + mock_context.get_provider_by_id.return_value = None + with pytest.raises(Exception, match="没有找到 ID 为"): + await ltm.get_image_caption("http://img.jpg", "missing-provider", "Describe") + + async def test_get_image_caption_raises_on_non_provider_type(self, ltm, mock_context): + mock_context.get_provider_by_id.return_value = "not_a_provider" + with pytest.raises(Exception, match="提供商类型错误"): + await ltm.get_image_caption("http://img.jpg", "bad-provider", "Describe") + + +class TestNeedActiveReply: + """Active reply decision logic.""" + + async def test_need_active_reply_false_when_disabled(self, ltm, mock_event): + with patch.object(LongTermMemory, "cfg", return_value={"enable_active_reply": False}): + assert await ltm.need_active_reply(mock_event) is False + + async def test_need_active_reply_false_for_private(self, ltm, mock_event): + mock_event.get_message_type.return_value = MessageType.FRIEND_MESSAGE + with patch.object(LongTermMemory, "cfg", return_value={"enable_active_reply": True, "ar_whitelist": [], "ar_method": "possibility_reply", "ar_possibility": 1.0}): + assert await ltm.need_active_reply(mock_event) is False + + async def test_need_active_reply_false_when_wake_command(self, ltm, mock_event): + mock_event.is_at_or_wake_command = True + with patch.object(LongTermMemory, "cfg", return_value={"enable_active_reply": True, "ar_whitelist": [], "ar_method": "possibility_reply", "ar_possibility": 1.0}): + assert await ltm.need_active_reply(mock_event) is False + + async def test_need_active_reply_whitelist_filters(self, ltm, mock_event): + with patch.object(LongTermMemory, "cfg", return_value={ + "enable_active_reply": True, + "ar_whitelist": ["other:group:999"], + "ar_method": "possibility_reply", + "ar_possibility": 1.0, + }): + assert await ltm.need_active_reply(mock_event) is False + + async def test_need_active_reply_possibility_triggers(self, ltm, mock_event): + with patch.object(LongTermMemory, "cfg", return_value={ + "enable_active_reply": True, + "ar_whitelist": [], + "ar_method": "possibility_reply", + "ar_possibility": 1.0, + }): + assert await ltm.need_active_reply(mock_event) is True + + +class TestHandleMessage: + """Message recording logic.""" + + @patch.object(LongTermMemory, "get_image_caption", return_value="sunset") + async def test_handle_message_plain_and_image(self, mock_get_caption, ltm, mock_event): + from astrbot.api.message_components import Plain, Image + + mock_event.get_messages.return_value = [ + Plain(text="Hello "), + Image(url="http://img.jpg"), + Plain(text="world"), + ] + mock_event.get_message_type.return_value = MessageType.GROUP_MESSAGE + with patch.object(LongTermMemory, "cfg", return_value={ + "max_cnt": 300, + "image_caption": True, + "image_caption_provider_id": "prov-1", + "image_caption_prompt": "Describe", + "enable_active_reply": False, + }): + await ltm.handle_message(mock_event) + assert len(ltm.session_chats["qq:group:123456"]) == 1 + msg = ltm.session_chats["qq:group:123456"][0] + assert "Hello" in msg + assert "sunset" in msg + assert "world" in msg + + async def test_handle_message_plain_only(self, ltm, mock_event): + from astrbot.api.message_components import Plain + + mock_event.get_messages.return_value = [ + Plain(text="How are you?"), + ] + mock_event.get_message_type.return_value = MessageType.GROUP_MESSAGE + with patch.object(LongTermMemory, "cfg", return_value={ + "max_cnt": 300, + "image_caption": False, + "enable_active_reply": False, + }): + await ltm.handle_message(mock_event) + msg = ltm.session_chats["qq:group:123456"][0] + assert "How are you?" in msg + + async def test_handle_message_ignores_private(self, ltm, mock_event): + mock_event.get_message_type.return_value = MessageType.FRIEND_MESSAGE + await ltm.handle_message(mock_event) + assert "qq:group:123456" not in ltm.session_chats + + async def test_handle_message_trims_excess(self, ltm, mock_event): + from astrbot.api.message_components import Plain + + mock_event.get_messages.return_value = [Plain(text="x")] + with patch.object(LongTermMemory, "cfg", return_value={ + "max_cnt": 2, + "image_caption": False, + "enable_active_reply": False, + }): + ltm.session_chats["qq:group:123456"] = ["old1", "old2"] + await ltm.handle_message(mock_event) + msgs = ltm.session_chats["qq:group:123456"] + assert len(msgs) == 2 + + +class TestOnReqLLM: + """LLM request modification.""" + + async def test_on_req_llm_skips_when_no_history(self, ltm, mock_event): + req = MagicMock() + await ltm.on_req_llm(mock_event, req) + req.assert_not_called() + + async def test_on_req_llm_adds_to_system_prompt(self, ltm, mock_event): + ltm.session_chats["qq:group:123456"] = ["msg1", "msg2"] + req = MagicMock() + req.system_prompt = "" + with patch.object(LongTermMemory, "cfg", return_value={"enable_active_reply": False}): + await ltm.on_req_llm(mock_event, req) + assert "msg1" in req.system_prompt + assert "msg2" in req.system_prompt + + async def test_on_req_llm_active_reply_overrides_prompt(self, ltm, mock_event): + ltm.session_chats["qq:group:123456"] = ["msg1"] + req = MagicMock() + req.prompt = "user query" + req.contexts = ["old_ctx"] + with patch.object(LongTermMemory, "cfg", return_value={"enable_active_reply": True}): + await ltm.on_req_llm(mock_event, req) + assert req.contexts == [] + assert "user query" in req.prompt + + +class TestAfterReqLLM: + """Post-LLM response recording.""" + + async def test_after_req_llm_records_response(self, ltm, mock_event): + ltm.session_chats["qq:group:123456"] = [] + resp = MagicMock() + resp.completion_text = "AI reply" + with patch.object(LongTermMemory, "cfg", return_value={"max_cnt": 300}): + await ltm.after_req_llm(mock_event, resp) + stored = ltm.session_chats["qq:group:123456"] + assert len(stored) == 1 + assert "AI reply" in stored[0] + + async def test_after_req_llm_skips_empty_response(self, ltm, mock_event): + ltm.session_chats["qq:group:123456"] = ["existing"] + resp = MagicMock() + resp.completion_text = None + await ltm.after_req_llm(mock_event, resp) + assert len(ltm.session_chats["qq:group:123456"]) == 1 + + async def test_after_req_llm_trims(self, ltm, mock_event): + ltm.session_chats["qq:group:123456"] = ["m1", "m2", "m3"] + resp = MagicMock() + resp.completion_text = "new" + with patch.object(LongTermMemory, "cfg", return_value={"max_cnt": 3}): + await ltm.after_req_llm(mock_event, resp) + assert len(ltm.session_chats["qq:group:123456"]) == 3 diff --git a/tests/unit/test_manager.py b/tests/unit/test_manager.py new file mode 100644 index 000000000..df45fe26d --- /dev/null +++ b/tests/unit/test_manager.py @@ -0,0 +1,335 @@ +"""Tests for astrbot.core.platform.manager — PlatformManager.""" + +import asyncio +from asyncio import Queue +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.platform.manager import PlatformManager +from astrbot.core.platform.platform import Platform, PlatformError, PlatformStatus + + +# --------------------------------------------------------------------------- +# A minimal concrete Platform we can wire into the manager without real I/O. +# --------------------------------------------------------------------------- +class DummyPlatform(Platform): + """Lightweight Platform used inside manager tests.""" + + def __init__( + self, + config: dict, + settings: dict, + event_queue: Queue, + **kwargs, + ) -> None: + super().__init__(config, event_queue) + self.settings = settings + self._terminated = False + + async def run(self) -> None: + pass + + def meta(self): + from astrbot.core.platform.platform_metadata import PlatformMetadata + return PlatformMetadata( + name=self.config.get("type", "dummy"), + description="dummy", + id=self.config.get("id", "dummy_id"), + ) + + async def terminate(self) -> None: + self._terminated = True + await super().terminate() + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def mock_config() -> MagicMock: + """A dict-like MagicMock that behaves like AstrBotConfig for key lookups.""" + cfg = MagicMock() + cfg.__getitem__.side_effect = lambda k: { + "platform": [], + "platform_settings": {"unique_session": True}, + }[k] + return cfg + + +@pytest.fixture +def event_queue() -> Queue: + return Queue() + + +@pytest.fixture +def manager(mock_config: MagicMock, event_queue: Queue) -> PlatformManager: + return PlatformManager(mock_config, event_queue) + + +# =================================================================== +# Construction +# =================================================================== + +class TestConstruction: + """PlatformManager.__init__ stores constructor arguments and initialises + empty collections.""" + + def test_stores_config(self, mock_config: MagicMock, event_queue: Queue): + m = PlatformManager(mock_config, event_queue) + assert m.astrbot_config is mock_config + assert m.event_queue is event_queue + + def test_platforms_config_and_settings_are_extracted( + self, mock_config: MagicMock, event_queue: Queue + ): + m = PlatformManager(mock_config, event_queue) + assert m.platforms_config == [] + assert m.settings == {"unique_session": True} + + def test_initial_collections_are_empty( + self, mock_config: MagicMock, event_queue: Queue + ): + m = PlatformManager(mock_config, event_queue) + assert m.platform_insts == [] + assert m._inst_map == {} + assert m._platform_tasks == {} + + +# =================================================================== +# _is_valid_platform_id +# =================================================================== + +class TestIsValidPlatformId: + """_is_valid_platform_id rejects None, empty, or ids containing ':'/'!'.""" + + def test_valid_id(self, manager: PlatformManager): + assert manager._is_valid_platform_id("my_platform") is True + assert manager._is_valid_platform_id("platform123") is True + + def test_none_is_invalid(self, manager: PlatformManager): + assert manager._is_valid_platform_id(None) is False + + def test_empty_string_is_invalid(self, manager: PlatformManager): + assert manager._is_valid_platform_id("") is False + + def test_colon_is_invalid(self, manager: PlatformManager): + assert manager._is_valid_platform_id("plat:form") is False + + def test_exclamation_is_invalid(self, manager: PlatformManager): + assert manager._is_valid_platform_id("plat!form") is False + + +# =================================================================== +# _sanitize_platform_id +# =================================================================== + +class TestSanitizePlatformId: + """_sanitize_platform_id replaces ':'/'!' with '_'.""" + + def test_clean_id_unchanged(self, manager: PlatformManager): + result, changed = manager._sanitize_platform_id("my_platform") + assert result == "my_platform" + assert changed is False + + def test_colon_replaced(self, manager: PlatformManager): + result, changed = manager._sanitize_platform_id("my:platform") + assert result == "my_platform" + assert changed is True + + def test_exclamation_replaced(self, manager: PlatformManager): + result, changed = manager._sanitize_platform_id("my!platform") + assert result == "my_platform" + assert changed is True + + def test_both_replaced(self, manager: PlatformManager): + result, changed = manager._sanitize_platform_id("a:b!c") + assert result == "a_b_c" + assert changed is True + + def test_none_returns_none_no_change(self, manager: PlatformManager): + result, changed = manager._sanitize_platform_id(None) + assert result is None + assert changed is False + + +# =================================================================== +# get_insts +# =================================================================== + +class TestGetInsts: + """get_insts returns the internal platform_insts list.""" + + def test_returns_same_list_object(self, manager: PlatformManager): + assert manager.get_insts() is manager.platform_insts + + def test_empty_by_default(self, manager: PlatformManager): + assert manager.get_insts() == [] + + +# =================================================================== +# get_all_stats +# =================================================================== + +class TestGetAllStats: + """get_all_stats aggregates stats from all registered platforms.""" + + def test_empty_when_no_platforms(self, manager: PlatformManager): + stats = manager.get_all_stats() + assert stats["platforms"] == [] + assert stats["summary"]["total"] == 0 + assert stats["summary"]["running"] == 0 + assert stats["summary"]["error"] == 0 + assert stats["summary"]["total_errors"] == 0 + + def test_aggregates_mixed_statuses(self, manager: PlatformManager): + run_mock = MagicMock(spec=Platform) + run_mock.get_stats.return_value = { + "id": "p1", + "type": "mock", + "display_name": "Mock1", + "status": PlatformStatus.RUNNING.value, + "started_at": None, + "error_count": 0, + "last_error": None, + "unified_webhook": False, + "meta": {}, + } + err_mock = MagicMock(spec=Platform) + err_mock.get_stats.return_value = { + "id": "p2", + "type": "mock", + "display_name": "Mock2", + "status": PlatformStatus.ERROR.value, + "started_at": None, + "error_count": 2, + "last_error": {"message": "fail"}, + "unified_webhook": False, + "meta": {}, + } + manager.platform_insts = [run_mock, err_mock] + + stats = manager.get_all_stats() + assert stats["summary"]["total"] == 2 + assert stats["summary"]["running"] == 1 + assert stats["summary"]["error"] == 1 + assert stats["summary"]["total_errors"] == 2 + + def test_recovers_from_broken_platform(self, manager: PlatformManager): + bad = MagicMock(spec=Platform) + bad.get_stats.side_effect = Exception("oops") + bad.config = {"id": "broken"} + manager.platform_insts = [bad] + + stats = manager.get_all_stats() + assert stats["summary"]["total"] == 1 + platform_info = stats["platforms"][0] + assert platform_info["id"] == "broken" + assert platform_info["status"] == "unknown" + assert platform_info["type"] == "unknown" + + +# =================================================================== +# terminate_platform +# =================================================================== + +class TestTerminatePlatform: + """terminate_platform removes the platform from maps and calls terminate.""" + + @pytest.mark.asyncio + async def test_removes_and_terminates(self, manager: PlatformManager): + p = DummyPlatform( + {"id": "test_id", "type": "dummy"}, {}, Queue() + ) + manager._inst_map["test_id"] = { + "inst": p, + "client_id": p.client_self_id, + } + manager.platform_insts.append(p) + + # Stop the underlying task machinery from actually creating tasks. + with patch.object(manager, "_stop_platform_task", new_callable=AsyncMock): + await manager.terminate_platform("test_id") + + assert p._terminated is True + assert "test_id" not in manager._inst_map + assert p not in manager.platform_insts + + +# =================================================================== +# terminate (all) +# =================================================================== + +class TestTerminateAll: + """terminate() cleans up all registered platforms.""" + + @pytest.mark.asyncio + async def test_terminates_all_platforms(self, manager: PlatformManager): + p1 = DummyPlatform({"id": "p1"}, {}, Queue()) + p2 = DummyPlatform({"id": "p2"}, {}, Queue()) + manager._inst_map["p1"] = {"inst": p1, "client_id": p1.client_self_id} + manager._inst_map["p2"] = {"inst": p2, "client_id": p2.client_self_id} + manager.platform_insts = [p1, p2] + + with patch.object(manager, "_stop_platform_task", new_callable=AsyncMock): + await manager.terminate() + + assert p1._terminated is True + assert p2._terminated is True + assert manager.platform_insts == [] + assert manager._inst_map == {} + assert manager._platform_tasks == {} + + +# =================================================================== +# load_platform +# =================================================================== + +class TestLoadPlatform: + """load_platform handles disabled flag and invalid IDs.""" + + @pytest.mark.asyncio + async def test_skips_disabled_platform(self, manager: PlatformManager): + """When enable is False, load_platform should return immediately.""" + with patch("astrbot.core.platform.manager.logger") as mock_logger: + await manager.load_platform({"enable": False}) + # info should not have been called with the loading message + for call in mock_logger.info.call_args_list: + assert "Loading" not in str(call) + + @pytest.mark.asyncio + async def test_sanitizes_invalid_platform_id( + self, manager: PlatformManager + ): + """A platform ID containing ':' should be sanitized automatically.""" + platform_cfg = { + "enable": True, + "type": "nonexistent_type", + "id": "bad:id", + } + with ( + patch.object(manager.astrbot_config, "save_config"), + patch("astrbot.core.platform.manager.logger"), + ): + await manager.load_platform(platform_cfg) + + assert platform_cfg["id"] == "bad_id" + + @pytest.mark.asyncio + async def test_logs_error_when_type_not_in_map( + self, manager: PlatformManager + ): + """If the platform type is not registered, load_platform logs an error.""" + platform_cfg = { + "enable": True, + "type": "no_such_adapter", + "id": "test_id", + } + with patch("astrbot.core.platform.manager.logger") as mock_logger: + await manager.load_platform(platform_cfg) + # Should have logged an error about adapter not found + error_messages = [ + str(c) for c in mock_logger.error.call_args_list + ] + assert any("not found" in msg for msg in error_messages) diff --git a/tests/unit/test_message_session.py b/tests/unit/test_message_session.py new file mode 100644 index 000000000..22f1b3dc0 --- /dev/null +++ b/tests/unit/test_message_session.py @@ -0,0 +1,169 @@ +"""Tests for astrbot.core.platform.message_session — MessageSession.""" + +import pytest + +from astrbot.core.platform.message_session import MessageSession, MessageSesion +from astrbot.core.platform.message_type import MessageType + + +# =================================================================== +# Construction +# =================================================================== + +class TestConstruction: + """MessageSession construction and post_init.""" + + def test_basic_construction(self): + session = MessageSession( + platform_name="test_platform", + message_type=MessageType.FRIEND_MESSAGE, + session_id="session_123", + ) + assert session.platform_name == "test_platform" + assert session.message_type == MessageType.FRIEND_MESSAGE + assert session.session_id == "session_123" + + def test_post_init_sets_platform_id_from_platform_name(self): + """platform_id should be auto-set to platform_name in __post_init__.""" + session = MessageSession( + platform_name="my_adapter", + message_type=MessageType.GROUP_MESSAGE, + session_id="g123", + ) + assert session.platform_id == "my_adapter" + + def test_platform_id_equals_platform_name(self): + """Confirm platform_id and platform_name hold the same value.""" + session = MessageSession( + platform_name="pname", + message_type=MessageType.OTHER_MESSAGE, + session_id="sid", + ) + assert session.platform_id == session.platform_name + + +# =================================================================== +# __str__ +# =================================================================== + +class TestStr: + """MessageSession.__str__ produces the unified-msg-origin format.""" + + def test_friend_message_format(self): + session = MessageSession( + platform_name="discord", + message_type=MessageType.FRIEND_MESSAGE, + session_id="user_42", + ) + assert str(session) == "discord:FriendMessage:user_42" + + def test_group_message_format(self): + session = MessageSession( + platform_name="telegram", + message_type=MessageType.GROUP_MESSAGE, + session_id="group_99", + ) + assert str(session) == "telegram:GroupMessage:group_99" + + def test_other_message_format(self): + session = MessageSession( + platform_name="system", + message_type=MessageType.OTHER_MESSAGE, + session_id="sys_1", + ) + assert str(session) == "system:OtherMessage:sys_1" + + +# =================================================================== +# from_str +# =================================================================== + +class TestFromStr: + """MessageSession.from_str parses the unified-msg-origin string.""" + + def test_parses_friend_message(self): + session = MessageSession.from_str("qq:FriendMessage:user_007") + assert session.platform_name == "qq" + assert session.platform_id == "qq" + assert session.message_type == MessageType.FRIEND_MESSAGE + assert session.session_id == "user_007" + + def test_parses_group_message(self): + session = MessageSession.from_str( + "slack:GroupMessage:channel_C01" + ) + assert session.platform_name == "slack" + assert session.message_type == MessageType.GROUP_MESSAGE + assert session.session_id == "channel_C01" + + @staticmethod + def test_roundtrip(): + """str -> from_str -> str should be lossless.""" + original = "wechat:FriendMessage:wx_abc" + session = MessageSession.from_str(original) + assert str(session) == original + + @staticmethod + def test_session_id_contains_colons(): + """If the session_id itself contains colons, only the first two split + boundaries are consumed; the rest is part of session_id.""" + raw = "platform:GroupMessage:user:name:extra" + session = MessageSession.from_str(raw) + assert session.platform_name == "platform" + assert session.message_type == MessageType.GROUP_MESSAGE + assert session.session_id == "user:name:extra" + assert str(session) == raw + + @staticmethod + def test_preserves_explicit_platform_id_after_roundtrip(): + """from_str sets platform_name from the parsed platform_id segment, + so platform_id == platform_name after a roundtrip.""" + session = MessageSession.from_str("my_id:FriendMessage:sid") + assert session.platform_name == "my_id" + assert session.platform_id == "my_id" + + +# =================================================================== +# Back-compat alias +# =================================================================== + +class TestAlias: + """MessageSesion (note the typo) should be an alias for MessageSession.""" + + def test_alias_is_same_class(self): + assert MessageSesion is MessageSession + + def test_alias_can_be_instantiated(self): + session = MessageSesion( + platform_name="alias_test", + message_type=MessageType.GROUP_MESSAGE, + session_id="alias_sid", + ) + assert isinstance(session, MessageSession) + assert session.platform_name == "alias_test" + + +# =================================================================== +# Dataclass equality +# =================================================================== + +class TestEquality: + """MessageSession is a dataclass so it inherits __eq__.""" + + def test_equal_sessions(self): + a = MessageSession( + platform_name="p", message_type=MessageType.FRIEND_MESSAGE, session_id="s" + ) + b = MessageSession( + platform_name="p", message_type=MessageType.FRIEND_MESSAGE, session_id="s" + ) + assert a == b + + def test_inequality_different_session_id(self): + a = MessageSession( + platform_name="p", message_type=MessageType.FRIEND_MESSAGE, session_id="s1" + ) + b = MessageSession( + platform_name="p", message_type=MessageType.FRIEND_MESSAGE, session_id="s2" + ) + assert a != b diff --git a/tests/unit/test_network_utils.py b/tests/unit/test_network_utils.py index ea3505e38..6a41097f8 100644 --- a/tests/unit/test_network_utils.py +++ b/tests/unit/test_network_utils.py @@ -1,8 +1,27 @@ -import ssl +"""Unit tests for astrbot.core.utils.network_utils. +Expands upon the existing 2 monkeypatch-based tests with comprehensive +coverage of is_connection_error, log_connection_failure, and +create_proxy_client edge cases. +""" + +import os +import ssl +from unittest.mock import MagicMock, patch + +import httpx import pytest -from astrbot.core.utils import network_utils +from astrbot.core.utils import network_utils as network_utils_module +from astrbot.core.utils.network_utils import ( + create_proxy_client, + is_connection_error, + log_connection_failure, +) + +# --------------------------------------------------------------------------- +# Existing tests (preserved verbatim) +# --------------------------------------------------------------------------- def test_create_proxy_client_reuses_shared_ssl_context( @@ -15,12 +34,12 @@ def test_create_proxy_client_reuses_shared_ssl_context( def __init__(self, **kwargs): captured_calls.append(kwargs) - monkeypatch.setattr(network_utils.httpx, "AsyncClient", _FakeAsyncClient) + monkeypatch.setattr(network_utils_module.httpx, "AsyncClient", _FakeAsyncClient) - network_utils.create_proxy_client("OpenAI") - network_utils.create_proxy_client("OpenAI", proxy="http://127.0.0.1:7890") - network_utils.create_proxy_client("OpenAI", headers=headers) - network_utils.create_proxy_client("OpenAI", proxy="") + network_utils_module.create_proxy_client("OpenAI") + network_utils_module.create_proxy_client("OpenAI", proxy="http://127.0.0.1:7890") + network_utils_module.create_proxy_client("OpenAI", headers=headers) + network_utils_module.create_proxy_client("OpenAI", proxy="") assert len(captured_calls) == 4 assert "proxy" not in captured_calls[0] @@ -43,9 +62,172 @@ def test_create_proxy_client_allows_verify_override( def __init__(self, **kwargs): captured_calls.append(kwargs) - monkeypatch.setattr(network_utils.httpx, "AsyncClient", _FakeAsyncClient) + monkeypatch.setattr(network_utils_module.httpx, "AsyncClient", _FakeAsyncClient) - network_utils.create_proxy_client("OpenAI", verify=custom_verify) + network_utils_module.create_proxy_client("OpenAI", verify=custom_verify) assert len(captured_calls) == 1 assert captured_calls[0]["verify"] is custom_verify + + +# --------------------------------------------------------------------------- +# is_connection_error +# --------------------------------------------------------------------------- + + +class TestIsConnectionError: + def test_httpx_connect_error(self): + assert is_connection_error(httpx.ConnectError("refused")) is True + + def test_httpx_timeout_errors(self): + for exc_cls in ( + httpx.ConnectTimeout, + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.PoolTimeout, + ): + assert is_connection_error(exc_cls("timeout")) is True + + def test_httpx_network_errors(self): + for exc_cls in (httpx.NetworkError, httpx.ProxyError, httpx.RequestError): + assert is_connection_error(exc_cls("err")) is True + + def test_builtin_network_exceptions(self): + for exc in (ConnectionError("conn"), TimeoutError("timeout"), OSError("os")): + assert is_connection_error(exc) is True + + def test_non_network_exceptions_return_false(self): + for exc in ( + ValueError("v"), + TypeError("t"), + RuntimeError("r"), + BaseException(), + ): + assert is_connection_error(exc) is False + + def test_httpx_non_network_errors_return_false(self): + request = httpx.Request("GET", "http://example.com") + response = httpx.Response(200, request=request) + for exc in ( + httpx.HTTPStatusError("err", request=request, response=response), + httpx.InvalidURL("invalid"), + httpx.CookieConflict("cookies"), + ): + assert is_connection_error(exc) is False + + def test_cause_chain_unwraps_to_network_error(self): + inner = httpx.ConnectError("inner") + outer = ValueError("wrapper") + outer.__cause__ = inner + assert is_connection_error(outer) is True + + def test_cause_chain_no_network_error(self): + inner = ValueError("inner") + outer = RuntimeError("outer") + outer.__cause__ = inner + assert is_connection_error(outer) is False + + def test_self_cause_does_not_infinite_loop(self): + exc = ValueError("self") + exc.__cause__ = exc + assert is_connection_error(exc) is False + + +# --------------------------------------------------------------------------- +# log_connection_failure +# --------------------------------------------------------------------------- + + +class TestLogConnectionFailure: + @patch("astrbot.core.utils.network_utils.logger.error") + def test_with_explicit_proxy(self, mock_log_error: MagicMock): + log_connection_failure("GPT", Exception("fail"), proxy="http://proxy:8080") + mock_log_error.assert_called_once() + msg = mock_log_error.call_args[0][0] + assert "GPT" in msg + assert "http://proxy:8080" in msg + assert "fail" in msg + + @patch("astrbot.core.utils.network_utils.logger.error") + def test_without_proxy_falls_to_simple_message(self, mock_log_error: MagicMock): + log_connection_failure("GPT", Exception("fail")) + mock_log_error.assert_called_once() + msg = mock_log_error.call_args[0][0] + assert "网络连接失败" in msg + assert "GPT" in msg + + @patch("astrbot.core.utils.network_utils.logger.error") + @patch.dict(os.environ, {"http_proxy": "http://env-proxy:3128"}, clear=False) + def test_falls_back_to_env_http_proxy( + self, mock_log_error: MagicMock + ): + log_connection_failure("GPT", Exception("fail")) + mock_log_error.assert_called_once() + msg = mock_log_error.call_args[0][0] + assert "http://env-proxy:3128" in msg + + @patch("astrbot.core.utils.network_utils.logger.error") + @patch.dict(os.environ, {"https_proxy": "https://secure-proxy:8443"}, clear=False) + def test_falls_back_to_env_https_proxy( + self, mock_log_error: MagicMock + ): + log_connection_failure("GPT", Exception("fail")) + mock_log_error.assert_called_once() + msg = mock_log_error.call_args[0][0] + assert "https://secure-proxy:8443" in msg + + @patch("astrbot.core.utils.network_utils.logger.error") + def test_empty_proxy_string_uses_no_proxy_message( + self, mock_log_error: MagicMock + ): + """Empty-string proxy should not trigger the proxy-specific log line.""" + with patch.dict(os.environ, {"http_proxy": "", "https_proxy": ""}): + log_connection_failure("GPT", Exception("fail"), proxy="") + mock_log_error.assert_called_once() + msg = mock_log_error.call_args[0][0] + assert "代理" not in msg + + @patch("astrbot.core.utils.network_utils.logger.error") + def test_error_type_name_in_message(self, mock_log_error: MagicMock): + log_connection_failure( + "GPT", httpx.ConnectTimeout("upstream timed out") + ) + mock_log_error.assert_called_once() + msg = mock_log_error.call_args[0][0] + assert "ConnectTimeout" in msg + + +# --------------------------------------------------------------------------- +# create_proxy_client +# --------------------------------------------------------------------------- + + +class TestCreateProxyClientIntegration: + """Integration-style smoke tests against the real httpx.AsyncClient.""" + + def test_returns_async_client_without_proxy(self): + client = create_proxy_client("Test") + assert isinstance(client, httpx.AsyncClient) + client.aclose() + + def test_returns_async_client_with_proxy(self): + client = create_proxy_client("Test", proxy="http://127.0.0.1:7890") + assert isinstance(client, httpx.AsyncClient) + client.aclose() + + def test_with_custom_headers(self): + headers = {"Authorization": "Bearer token"} + client = create_proxy_client("Test", headers=headers) + assert isinstance(client, httpx.AsyncClient) + client.aclose() + + def test_with_explicit_ssl_context(self): + ctx = ssl.create_default_context() + client = create_proxy_client("Test", verify=ctx) + assert isinstance(client, httpx.AsyncClient) + client.aclose() + + def test_with_verify_disabled(self): + client = create_proxy_client("Test", verify=False) + assert isinstance(client, httpx.AsyncClient) + client.aclose() diff --git a/tests/unit/test_pipeline_bootstrap.py b/tests/unit/test_pipeline_bootstrap.py new file mode 100644 index 000000000..a83dceedd --- /dev/null +++ b/tests/unit/test_pipeline_bootstrap.py @@ -0,0 +1,238 @@ +"""Unit tests for astrbot.core.pipeline.bootstrap module.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, call, patch + +import pytest + +from astrbot.core.pipeline.bootstrap import ( + _BUILTIN_STAGE_MODULES, + _EXPECTED_STAGE_NAMES, + _builtin_stages_registered, + ensure_builtin_stages_registered, +) + + +class TestEnsureBuiltinStagesRegistered: + """Tests for ensure_builtin_stages_registered().""" + + def teardown_method(self): + """Reset global state after each test.""" + # Use import and direct assignment to reset + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_already_registered_flag_short_circuits(self): + """When _builtin_stages_registered is True, return immediately.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = True + + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + + mock_import.assert_not_called() + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_all_expected_stages_present_skips_import(self): + """When all expected stages are already in registered_stages, skip importing.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + + # Populate registered_stages with all expected stage names + for name in _EXPECTED_STAGE_NAMES: + cls = MagicMock() + cls.__name__ = name + bootstrap_mod.registered_stages.append(cls) + + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + + mock_import.assert_not_called() + assert bootstrap_mod._builtin_stages_registered is True + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_registers_missing_stages(self): + """When stages are missing, import built-in modules.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + + expected_calls = [call(mod) for mod in _BUILTIN_STAGE_MODULES] + mock_import.assert_has_calls(expected_calls, any_order=True) + assert mock_import.call_count == len(_BUILTIN_STAGE_MODULES) + assert bootstrap_mod._builtin_stages_registered is True + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_idempotent(self): + """Calling ensure_builtin_stages_registered twice only imports once.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + ensure_builtin_stages_registered() + + # Should only import on first call + mock_import.assert_called_once() + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_global_flag_set_after_import(self): + """Verify the global flag is set after a full registration.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + with patch("astrbot.core.pipeline.bootstrap.import_module"): + ensure_builtin_stages_registered() + + assert bootstrap_mod._builtin_stages_registered is True + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_already_registered_flag_persists(self): + """Verify that once set, the flag persists across calls.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + with patch("astrbot.core.pipeline.bootstrap.import_module"): + ensure_builtin_stages_registered() + + # Call again with import_module mocked to raise if called + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + + mock_import.assert_not_called() + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_partial_stages_present_still_imports(self): + """When only some expected stages are present, still import all modules.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + # Add only one of the expected stages + cls = MagicMock() + cls.__name__ = "ProcessStage" + bootstrap_mod.registered_stages.append(cls) + + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + + expected_calls = [call(mod) for mod in _BUILTIN_STAGE_MODULES] + mock_import.assert_has_calls(expected_calls, any_order=True) + assert mock_import.call_count == len(_BUILTIN_STAGE_MODULES) + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_stage_names_check_exact(self): + """Verify the check uses __name__ comparison, not identity.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + # Add stages with correct names + for name in _EXPECTED_STAGE_NAMES: + cls = MagicMock() + cls.__name__ = name + bootstrap_mod.registered_stages.append(cls) + + with patch( + "astrbot.core.pipeline.bootstrap.import_module", + ) as mock_import: + ensure_builtin_stages_registered() + + mock_import.assert_not_called() + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_expected_stage_names_are_correct(self): + """Verify _EXPECTED_STAGE_NAMES matches the expected set.""" + expected = { + "WakingCheckStage", + "WhitelistCheckStage", + "SessionStatusCheckStage", + "RateLimitStage", + "ContentSafetyCheckStage", + "PreProcessStage", + "ProcessStage", + "ResultDecorateStage", + "RespondStage", + } + assert _EXPECTED_STAGE_NAMES == expected + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_builtin_stage_modules_count(self): + """Verify the number of builtin stage modules matches expected.""" + assert len(_BUILTIN_STAGE_MODULES) == 9 + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_all_builtin_modules_under_pipeline(self): + """Verify all builtin modules are under astrbot.core.pipeline.""" + for mod in _BUILTIN_STAGE_MODULES: + assert mod.startswith("astrbot.core.pipeline.") + assert mod.endswith(".stage") + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_module_paths_alignment(self): + """Verify each expected stage name has a corresponding module path.""" + from astrbot.core.pipeline.bootstrap import _BUILTIN_STAGE_MODULES, _EXPECTED_STAGE_NAMES + + # Derive expected names from module paths + derived_names = set() + for mod_path in _BUILTIN_STAGE_MODULES: + parts = mod_path.split(".") + # For modules like astrbot.core.pipeline.process_stage.stage + # the stage name is found in the penultimate segment + if parts[-2] in ("process_stage", "respond"): + # Special cases: ProcessStage, RespondStage + if parts[-2] == "process_stage": + derived_names.add("ProcessStage") + elif parts[-2] == "respond": + derived_names.add("RespondStage") + else: + derived_names.add(f"{parts[-2].title().replace('_', '')}Stage") + else: + name = parts[-2].replace("_", " ").title().replace(" ", "") + name += "Stage" + derived_names.add(name) + + assert _EXPECTED_STAGE_NAMES == derived_names + + @patch("astrbot.core.pipeline.bootstrap.registered_stages", new=[]) + def test_real_registration_smoke(self): + """Smoke test: calling ensure_builtin_stages_registered with actual modules.""" + import astrbot.core.pipeline.bootstrap as bootstrap_mod + + bootstrap_mod._builtin_stages_registered = False + bootstrap_mod.registered_stages.clear() + + # This should import the actual modules + ensure_builtin_stages_registered() + + stage_names = {cls.__name__ for cls in bootstrap_mod.registered_stages} + assert _EXPECTED_STAGE_NAMES.issubset(stage_names) diff --git a/tests/unit/test_pipeline_process_stage.py b/tests/unit/test_pipeline_process_stage.py new file mode 100644 index 000000000..98e0c16d1 --- /dev/null +++ b/tests/unit/test_pipeline_process_stage.py @@ -0,0 +1,346 @@ +"""Unit tests for astrbot.core.pipeline.process_stage.stage.ProcessStage.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.pipeline.process_stage.stage import ProcessStage +from astrbot.core.provider.entities import ProviderRequest + + +@pytest.fixture +def mock_context(): + """Create a mock PipelineContext.""" + ctx = MagicMock() + ctx.astrbot_config = { + "provider_settings": {"enable": True, "wake_prefix": ""}, + "wake_prefix": ["bot"], + } + ctx.plugin_manager = MagicMock() + ctx.plugin_manager.context = MagicMock() + return ctx + + +@pytest.fixture +def mock_event(): + """Create a mock AstrMessageEvent.""" + event = MagicMock() + event.get_extra.return_value = None + event.is_stopped.return_value = False + event._has_send_oper = False + event.is_at_or_wake_command = True + event.call_llm = True + event.set_extra = MagicMock() + return event + + +@pytest.fixture +def stage(mock_context): + """Create a ProcessStage with mocked sub-stages.""" + stage = ProcessStage() + + stage.agent_sub_stage = MagicMock() + stage.agent_sub_stage.process = AsyncMock() + + stage.star_request_sub_stage = MagicMock() + stage.star_request_sub_stage.process = AsyncMock() + + stage.ctx = mock_context + stage.config = mock_context.astrbot_config + stage.plugin_manager = mock_context.plugin_manager + stage.sdk_plugin_bridge = None + return stage + + +class TestProcessStageInitialize: + """Tests for ProcessStage.initialize().""" + + @pytest.mark.asyncio + async def test_initialize_sets_attributes(self, mock_context): + """Verify initialize sets ctx, config, plugin_manager.""" + stage = ProcessStage() + with ( + patch.object(stage, "agent_sub_stage") as mock_agent, + patch.object(stage, "star_request_sub_stage") as mock_star, + ): + await stage.initialize(mock_context) + + assert stage.ctx is mock_context + assert stage.config is mock_context.astrbot_config + assert stage.plugin_manager is mock_context.plugin_manager + + @pytest.mark.asyncio + async def test_initialize_creates_sub_stages(self, mock_context): + """Verify initialize creates and initializes sub-stages.""" + stage = ProcessStage() + await stage.initialize(mock_context) + + assert hasattr(stage, "agent_sub_stage") + assert hasattr(stage, "star_request_sub_stage") + assert stage.agent_sub_stage is not None + assert stage.star_request_sub_stage is not None + + @pytest.mark.asyncio + async def test_initialize_sdk_plugin_bridge_present(self, mock_context): + """Verify sdk_plugin_bridge is set when plugin_manager has it.""" + mock_context.plugin_manager.context.sdk_plugin_bridge = MagicMock() + stage = ProcessStage() + await stage.initialize(mock_context) + + assert stage.sdk_plugin_bridge is mock_context.plugin_manager.context.sdk_plugin_bridge + + @pytest.mark.asyncio + async def test_initialize_sdk_plugin_bridge_absent(self, mock_context): + """Verify sdk_plugin_bridge is None when plugin_manager has no context attr.""" + mock_context.plugin_manager = MagicMock(spec=[]) # no context attribute + stage = ProcessStage() + await stage.initialize(mock_context) + + assert stage.sdk_plugin_bridge is None + + +class TestProcessStageProcess: + """Tests for ProcessStage.process().""" + + async def _collect(self, async_gen): + """Helper to collect all items from an async generator.""" + results = [] + async for item in async_gen: + results.append(item) + return results + + @pytest.mark.asyncio + async def test_process_no_activated_handlers_and_no_sdk_bridge( + self, stage, mock_event, + ): + """When activated_handlers is None and no sdk bridge, skip handler path.""" + mock_event.get_extra.return_value = None + stage.sdk_plugin_bridge = None + + results = await self._collect(stage.process(mock_event)) + + assert results == [] # provider enabled, should still reach agent_sub_stage + stage.agent_sub_stage.process.assert_awaited_once() + + @pytest.mark.asyncio + async def test_process_activated_handlers_with_provider_request( + self, stage, mock_event, + ): + """When star_request yields a ProviderRequest, agent_sub_stage is called.""" + pr = ProviderRequest(prompt="hello") + + async def star_gen(_event): + yield pr + + stage.star_request_sub_stage.process = star_gen + + agent_called = False + + async def agent_gen(_event): + nonlocal agent_called + agent_called = True + yield None + + stage.agent_sub_stage.process = agent_gen + + results = await self._collect(stage.process(mock_event)) + + assert agent_called is True + mock_event.set_extra.assert_any_call("provider_request", pr) + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_activated_handlers_with_non_provider_request( + self, stage, mock_event, + ): + """When star_request yields a non-ProviderRequest, yield directly.""" + async def star_gen(_event): + yield "some_other_result" + + stage.star_request_sub_stage.process = star_gen + stage.sdk_plugin_bridge = None + + results = await self._collect(stage.process(mock_event)) + + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_activated_handlers_provider_request_empty_agent( + self, stage, mock_event, + ): + """When agent_sub_stage yields nothing, should still yield once.""" + pr = ProviderRequest(prompt="hi") + + async def star_gen(_event): + yield pr + + stage.star_request_sub_stage.process = star_gen + + async def empty_agent_gen(_event): + if False: + yield None + + stage.agent_sub_stage.process = empty_agent_gen + + results = await self._collect(stage.process(mock_event)) + + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_sdk_plugin_bridge_sent_message(self, stage, mock_event): + """When sdk_plugin_bridge.dispatch_message returns sent_message=True.""" + mock_bridge = MagicMock() + mock_result = MagicMock() + mock_result.sent_message = True + mock_result.stopped = False + mock_bridge.dispatch_message = AsyncMock(return_value=mock_result) + stage.sdk_plugin_bridge = mock_bridge + + results = await self._collect(stage.process(mock_event)) + + # Should have yielded due to sent_message + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_sdk_plugin_bridge_stopped(self, stage, mock_event): + """When sdk_plugin_bridge.dispatch_message returns stopped=True.""" + mock_bridge = MagicMock() + mock_result = MagicMock() + mock_result.sent_message = False + mock_result.stopped = True + mock_bridge.dispatch_message = AsyncMock(return_value=mock_result) + stage.sdk_plugin_bridge = mock_bridge + + results = await self._collect(stage.process(mock_event)) + + # Should have yielded due to stopped + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_sdk_plugin_bridge_none_and_event_has_send_oper( + self, stage, mock_event, + ): + """When sdk bridge is None and _has_send_oper is True, skip LLM call.""" + stage.sdk_plugin_bridge = None + mock_event._has_send_oper = True + + results = await self._collect(stage.process(mock_event)) + + # Should NOT call agent_sub_stage for LLM + assert results == [] + + @pytest.mark.asyncio + async def test_process_provider_disabled(self, stage, mock_event): + """When provider_settings enable is False, return early.""" + stage.ctx.astrbot_config["provider_settings"]["enable"] = False + stage.sdk_plugin_bridge = None + mock_event.get_extra.return_value = None + + results = await self._collect(stage.process(mock_event)) + + assert results == [] + + @pytest.mark.asyncio + async def test_process_llm_triggered(self, stage, mock_event): + """When all conditions met, agent_sub_stage is called for LLM.""" + stage.sdk_plugin_bridge = None + agent_called = False + + async def agent_gen(_event): + nonlocal agent_called + agent_called = True + yield None + + stage.agent_sub_stage.process = agent_gen + + results = await self._collect(stage.process(mock_event)) + + assert agent_called is True + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_llm_skipped_not_at_wake(self, stage, mock_event): + """When is_at_or_wake_command is False, skip LLM call.""" + stage.sdk_plugin_bridge = None + mock_event.is_at_or_wake_command = False + + results = await self._collect(stage.process(mock_event)) + + assert results == [] + + @pytest.mark.asyncio + async def test_process_llm_skipped_event_stopped_with_result( + self, stage, mock_event, + ): + """When event is stopped and effective_result exists, skip LLM call.""" + stage.sdk_plugin_bridge = None + mock_event.is_stopped.return_value = True + + mock_bridge = MagicMock() + mock_bridge.get_effective_should_call_llm.return_value = True + mock_bridge.get_effective_result.return_value = "some_result" + stage.sdk_plugin_bridge = mock_bridge + + results = await self._collect(stage.process(mock_event)) + + # Event is stopped with result, skip LLM + assert results == [] + + @pytest.mark.asyncio + async def test_process_sdk_bridge_with_effective_methods( + self, stage, mock_event, + ): + """When sdk_plugin_bridge has get_effective_* methods, they are used.""" + mock_bridge = MagicMock() + mock_bridge.get_effective_should_call_llm.return_value = False + stage.sdk_plugin_bridge = mock_bridge + + await self._collect(stage.process(mock_event)) + + mock_bridge.get_effective_should_call_llm.assert_called_once_with( + mock_event, + ) + + @pytest.mark.asyncio + async def test_process_with_activated_handlers_and_stopped( + self, stage, mock_event, + ): + """When event is stopped after star_request processing, stop propagation.""" + async def star_gen(_event): + yield ProviderRequest(prompt="test") + + stage.star_request_sub_stage.process = star_gen + + async def agent_gen(_event): + yield None + + stage.agent_sub_stage.process = agent_gen + + mock_event.is_stopped.return_value = True + + results = await self._collect(stage.process(mock_event)) + + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_process_event_call_llm_false(self, stage, mock_event): + """When event.call_llm is False and no sdk bridge, should_call_llm is True (inverted).""" + stage.sdk_plugin_bridge = None + mock_event.call_llm = False + + agent_called = False + + async def agent_gen(_event): + nonlocal agent_called + agent_called = True + yield None + + stage.agent_sub_stage.process = agent_gen + + results = await self._collect(stage.process(mock_event)) + + # event.call_llm is False, and no sdk_bridge -> should_call_llm = not False = True + assert agent_called is True + assert len(results) >= 1 diff --git a/tests/unit/test_pipeline_respond_stage.py b/tests/unit/test_pipeline_respond_stage.py new file mode 100644 index 000000000..e7f00b215 --- /dev/null +++ b/tests/unit/test_pipeline_respond_stage.py @@ -0,0 +1,530 @@ +"""Unit tests for astrbot.core.pipeline.respond.stage.RespondStage.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.message.components import ( + At, + ComponentType, + Face, + File, + Forward, + Image, + Plain, + Record, + Reply, +) +from astrbot.core.message.message_event_result import ( + MessageEventResult, + ResultContentType, +) +from astrbot.core.pipeline.respond.stage import RespondStage + + +@pytest.fixture +def mock_config(): + """Create a mock AstrBotConfig for RespondStage initialization.""" + return { + "platform_settings": { + "reply_with_mention": True, + "reply_with_quote": False, + "segmented_reply": { + "enable": False, + "only_llm_result": False, + "interval_method": "random", + "log_base": 2, + "interval": "1.5, 3.5", + }, + "path_mapping": [], + }, + "provider_settings": { + "enable": True, + "unsupported_streaming_strategy": "realtime_segmenting", + }, + } + + +@pytest.fixture +def mock_context(mock_config): + """Create a mock PipelineContext.""" + ctx = MagicMock() + ctx.astrbot_config = mock_config + return ctx + + +@pytest.fixture +def mock_event(): + """Create a mock AstrMessageEvent.""" + event = MagicMock() + event.get_extra.return_value = False + event.get_platform_name.return_value = "qq" + event.get_sender_name.return_value = "test_user" + event.get_sender_id.return_value = "12345" + event.get_platform_id.return_value = "platform_1" + return event + + +@pytest.fixture +def stage(mock_context): + """Create an initialized RespondStage.""" + stage = RespondStage() + return stage + + +class TestRespondStageInitialize: + """Tests for RespondStage.initialize().""" + + @pytest.mark.asyncio + async def test_initialize_sets_attributes(self, stage, mock_context): + """Verify initialize reads config and sets attributes.""" + await stage.initialize(mock_context) + + assert stage.ctx is mock_context + assert stage.config is mock_context.astrbot_config + assert stage.reply_with_mention is True + assert stage.reply_with_quote is False + assert stage.enable_seg is False + assert stage.only_llm_result is False + assert stage.interval_method == "random" + assert stage.log_base == 2 + assert stage.interval == [1.5, 3.5] + + @pytest.mark.asyncio + async def test_initialize_segmented_enabled(self, mock_context): + """Verify segmented reply config is parsed when enabled.""" + mock_context.astrbot_config["platform_settings"]["segmented_reply"][ + "enable" + ] = True + stage = RespondStage() + await stage.initialize(mock_context) + + assert stage.enable_seg is True + assert stage.interval == [1.5, 3.5] + + @pytest.mark.asyncio + async def test_initialize_segmented_invalid_interval(self, mock_context): + """Verify invalid interval string falls back gracefully.""" + mock_context.astrbot_config["platform_settings"]["segmented_reply"][ + "enable" + ] = True + mock_context.astrbot_config["platform_settings"]["segmented_reply"][ + "interval" + ] = "not_a_number" + stage = RespondStage() + await stage.initialize(mock_context) + + # Should fall back to [1.5, 3.5] + assert stage.interval == [1.5, 3.5] + + +class TestRespondStageWordCnt: + """Tests for RespondStage._word_cnt().""" + + @pytest.mark.asyncio + async def test_word_cnt_ascii(self, stage): + """Verify ASCII text is split and counted.""" + count = await stage._word_cnt("hello world foo bar") + assert count == 4 + + @pytest.mark.asyncio + async def test_word_cnt_ascii_single(self, stage): + """Verify single ASCII word.""" + count = await stage._word_cnt("hello") + assert count == 1 + + @pytest.mark.asyncio + async def test_word_cnt_non_ascii(self, stage): + """Verify non-ASCII text counts alphanumeric characters.""" + count = await stage._word_cnt("你好世界") + assert count == 4 + + @pytest.mark.asyncio + async def test_word_cnt_mixed(self, stage): + """Verify mixed text counts all alnum characters.""" + count = await stage._word_cnt("hello你好world世界") + assert count == 18 # hello(5) + 你好(4) + world(5) + 世界(4) = 18 + + @pytest.mark.asyncio + async def test_word_cnt_empty(self, stage): + """Verify empty string returns 0.""" + count = await stage._word_cnt("") + assert count == 0 + + +class TestRespondStageCalcCompInterval: + """Tests for RespondStage._calc_comp_interval().""" + + @pytest.mark.asyncio + async def test_calc_interval_log_plain(self, stage): + """Verify log interval for Plain component.""" + stage.interval_method = "log" + stage.log_base = 2 + plain = Plain(text="hello world") + + with patch("astrbot.core.pipeline.respond.stage.random.uniform") as mock_uniform: + mock_uniform.return_value = 2.0 + interval = await stage._calc_comp_interval(plain) + + assert interval == 2.0 + mock_uniform.assert_called_once() + + @pytest.mark.asyncio + async def test_calc_interval_log_non_plain(self, stage): + """Verify log interval for non-Plain component uses 1-1.75 range.""" + stage.interval_method = "log" + image = Image(file="/path/to/img.jpg") + + with patch("astrbot.core.pipeline.respond.stage.random.uniform") as mock_uniform: + mock_uniform.return_value = 1.5 + interval = await stage._calc_comp_interval(image) + + assert interval == 1.5 + mock_uniform.assert_called_once_with(1, 1.75) + + @pytest.mark.asyncio + async def test_calc_interval_random(self, stage): + """Verify random interval uses configured interval range.""" + stage.interval_method = "random" + stage.interval = [2.0, 4.0] + plain = Plain(text="test") + + with patch("astrbot.core.pipeline.respond.stage.random.uniform") as mock_uniform: + mock_uniform.return_value = 3.0 + interval = await stage._calc_comp_interval(plain) + + assert interval == 3.0 + mock_uniform.assert_called_once_with(2.0, 4.0) + + +class TestRespondStageHasMeaningfulContent: + """Tests for RespondStage._has_meaningful_content().""" + + def test_plain_with_text(self, stage): + """Verify Plain with text returns True.""" + comp = Plain(text="hello") + assert stage._has_meaningful_content(comp) is True + + def test_plain_whitespace_text(self, stage): + """Verify Plain with whitespace returns False.""" + comp = Plain(text=" ") + assert stage._has_meaningful_content(comp) is False + + def test_image_with_url(self, stage): + """Verify Image with url returns True.""" + comp = Image(file="http://example.com/img.jpg") + assert stage._has_meaningful_content(comp) is True + + def test_image_with_file_id(self, stage): + """Verify Image with file_id returns True.""" + comp = Image(file_id="abc123") + assert stage._has_meaningful_content(comp) is True + + def test_image_empty(self, stage): + """Verify Image without url or file_id returns False.""" + comp = Image() + assert stage._has_meaningful_content(comp) is False + + def test_face_with_id(self, stage): + """Verify Face with id returns True.""" + comp = Face(id=123) + assert stage._has_meaningful_content(comp) is True + + def test_face_no_id(self, stage): + """Verify Face without id returns False.""" + comp = Face(id=None) + assert stage._has_meaningful_content(comp) is False + + def test_at_with_qq(self, stage): + """Verify At with qq returns True.""" + comp = At(qq="12345") + assert stage._has_meaningful_content(comp) is True + + def test_at_no_qq(self, stage): + """Verify At without qq returns False.""" + comp = At(qq=None) + assert stage._has_meaningful_content(comp) is False + + def test_reply_with_id(self, stage): + """Verify Reply with id returns True.""" + comp = Reply(id="abc", sender_id="user1") + assert stage._has_meaningful_content(comp) is True + + def test_reply_no_id(self, stage): + """Verify Reply without id returns False.""" + comp = Reply(id=None, sender_id="user1") + assert stage._has_meaningful_content(comp) is False + + def test_forward_with_id(self, stage): + """Verify Forward with id returns True.""" + comp = Forward(id="abc123") + assert stage._has_meaningful_content(comp) is True + + +class TestRespondStageIsEmptyMessageChain: + """Tests for RespondStage._is_empty_message_chain().""" + + @pytest.mark.asyncio + async def test_empty_list(self, stage): + """Verify empty list returns True.""" + assert await stage._is_empty_message_chain([]) is True + + @pytest.mark.asyncio + async def test_chain_with_meaningful_content(self, stage): + """Verify chain with valid content returns False.""" + chain = [Plain(text="hello")] + assert await stage._is_empty_message_chain(chain) is False + + @pytest.mark.asyncio + async def test_chain_all_empty(self, stage): + """Verify chain with all empty components returns True.""" + chain = [Plain(text="")] + assert await stage._is_empty_message_chain(chain) is True + + @pytest.mark.asyncio + async def test_chain_mixed_empty_and_valid(self, stage): + """Verify chain with mix of empty and valid returns False.""" + chain = [Plain(text=""), Plain(text="hello")] + assert await stage._is_empty_message_chain(chain) is False + + +class TestRespondStageIsSegReplyRequired: + """Tests for RespondStage.is_seg_reply_required().""" + + def test_seg_disabled(self, stage): + """Verify returns False when segmented reply is disabled.""" + stage.enable_seg = False + event = MagicMock() + assert stage.is_seg_reply_required(event) is False + + def test_seg_no_result(self, stage): + """Verify returns False when event has no result.""" + stage.enable_seg = True + event = MagicMock() + event.get_result.return_value = None + assert stage.is_seg_reply_required(event) is False + + def test_seg_only_llm_not_model(self, stage): + """Verify returns False when only_llm_result is True and result is not model.""" + stage.enable_seg = True + stage.only_llm_result = True + event = MagicMock() + result = MagicMock() + result.is_model_result.return_value = False + event.get_result.return_value = result + assert stage.is_seg_reply_required(event) is False + + def test_seg_excluded_platform(self, stage): + """Verify returns False for excluded platforms.""" + stage.enable_seg = True + stage.only_llm_result = False + event = MagicMock() + result = MagicMock() + result.is_model_result.return_value = True + event.get_result.return_value = result + event.get_platform_name.return_value = "qq_official" + assert stage.is_seg_reply_required(event) is False + + def test_seg_all_conditions_met(self, stage): + """Verify returns True when all conditions are met.""" + stage.enable_seg = True + stage.only_llm_result = False + event = MagicMock() + result = MagicMock() + result.is_model_result.return_value = True + event.get_result.return_value = result + event.get_platform_name.return_value = "qq" + assert stage.is_seg_reply_required(event) is True + + +class TestRespondStageExtractComp: + """Tests for RespondStage._extract_comp().""" + + def test_extract_with_modify(self, stage): + """Verify extraction removes extracted types from original list.""" + raw_chain = [ + Plain(text="hello"), + At(qq="12345"), + Plain(text="world"), + Reply(id="r1", sender_id="u1"), + ] + extracted = stage._extract_comp( + raw_chain, + {ComponentType.At, ComponentType.Reply}, + modify_raw_chain=True, + ) + + assert len(extracted) == 2 + assert all(c.type in {ComponentType.At, ComponentType.Reply} for c in extracted) + assert len(raw_chain) == 2 + assert all(c.type == ComponentType.Plain for c in raw_chain) + + def test_extract_without_modify(self, stage): + """Verify extraction does not modify original chain.""" + raw_chain = [ + Plain(text="hello"), + At(qq="12345"), + ] + original_len = len(raw_chain) + extracted = stage._extract_comp( + raw_chain, + {ComponentType.At}, + modify_raw_chain=False, + ) + + assert len(extracted) == 1 + assert extracted[0].type == ComponentType.At + assert len(raw_chain) == original_len # unchanged + + def test_extract_no_match(self, stage): + """Verify extraction returns empty when no types match.""" + raw_chain = [Plain(text="hello")] + extracted = stage._extract_comp( + raw_chain, + {ComponentType.Image}, + modify_raw_chain=False, + ) + assert extracted == [] + + +class TestRespondStageProcess: + """Tests for RespondStage.process().""" + + @pytest.mark.asyncio + async def test_process_result_none(self, stage, mock_event): + """Verify process returns when result is None.""" + mock_event.get_result.return_value = None + await stage.process(mock_event) + # Should return without sending + mock_event.send.assert_not_called() + + @pytest.mark.asyncio + async def test_process_streaming_finished(self, stage, mock_event): + """Verify process returns when streaming is finished.""" + result = MagicMock() + result.result_content_type = ResultContentType.STREAMING_FINISH + mock_event.get_result.return_value = result + await stage.process(mock_event) + + mock_event.set_extra.assert_called_once_with("_streaming_finished", True) + + @pytest.mark.asyncio + async def test_process_streaming_finish_prevents_duplicate_send( + self, stage, mock_event, + ): + """Verify prevent duplicate send after streaming finish.""" + mock_event.get_extra.return_value = True # _streaming_finished already True + result = MagicMock() + result.result_content_type = ResultContentType.GENERAL_RESULT + mock_event.get_result.return_value = result + await stage.process(mock_event) + + mock_event.send.assert_not_called() + + @pytest.mark.asyncio + async def test_process_streaming_result(self, stage, mock_event): + """Verify STREAMING_RESULT is delivered directly to event.send_streaming.""" + async def dummy_async_stream(): + yield "chunk1" + + result = MagicMock() + result.result_content_type = ResultContentType.STREAMING_RESULT + result.async_stream = dummy_async_stream() + mock_event.get_result.return_value = result + + await stage.process(mock_event) + + mock_event.send_streaming.assert_called_once() + + @pytest.mark.asyncio + async def test_process_normal_chain(self, stage, mock_event): + """Verify normal message chain is sent via event.send.""" + result = MessageEventResult() + result.chain = [Plain(text="hello")] + mock_event.get_result.return_value = result + + await stage.process(mock_event) + + mock_event.send.assert_called_once() + + @pytest.mark.asyncio + async def test_process_empty_chain(self, stage, mock_event): + """Verify empty chain skips sending.""" + result = MessageEventResult() + result.chain = [Plain(text="")] + mock_event.get_result.return_value = result + + await stage.process(mock_event) + + mock_event.send.assert_not_called() + + @pytest.mark.asyncio + async def test_process_chain_with_path_mapping(self, stage, mock_event): + """Verify path mapping is applied to File components.""" + stage.platform_settings = { + "path_mapping": [["/old/path", "/new/path"]], + } + + with patch( + "astrbot.core.pipeline.respond.stage.path_Mapping", + return_value="/new/path/file.txt", + ): + result = MessageEventResult() + result.chain = [File(file="/old/path/file.txt")] + mock_event.get_result.return_value = result + + await stage.process(mock_event) + + mock_event.send.assert_called_once() + + @pytest.mark.asyncio + async def test_process_chain_with_record_forced_separate( + self, stage, mock_event, + ): + """Verify Record components are sent separately.""" + result = MessageEventResult() + result.chain = [Record(file="audio.mp3"), Plain(text="hello")] + mock_event.get_result.return_value = result + + await stage.process(mock_event) + + # Record should be sent separately, then the rest + assert mock_event.send.call_count >= 1 + + @pytest.mark.asyncio + async def test_process_segmented_reply(self, stage, mock_event): + """Verify segmented reply sends each component with delay.""" + stage.enable_seg = True + stage.only_llm_result = False + stage.interval_method = "random" + stage.interval = [0.1, 0.2] + + result = MessageEventResult() + result.chain = [Plain(text="part1"), Plain(text="part2")] + mock_event.get_result.return_value = result + mock_event.get_platform_name.return_value = "qq" + + with patch( + "astrbot.core.pipeline.respond.stage.asyncio.sleep", + AsyncMock(), + ): + await stage.process(mock_event) + + assert mock_event.send.call_count >= 1 + + @pytest.mark.asyncio + async def test_process_event_hook_called(self, stage, mock_event): + """Verify call_event_hook is invoked after sending.""" + result = MessageEventResult() + result.chain = [Plain(text="hello")] + mock_event.get_result.return_value = result + + with patch( + "astrbot.core.pipeline.respond.stage.call_event_hook", + AsyncMock(return_value=False), + ): + await stage.process(mock_event) + + mock_event.clear_result.assert_called_once() diff --git a/tests/unit/test_pipeline_scheduler.py b/tests/unit/test_pipeline_scheduler.py new file mode 100644 index 000000000..db13010d1 --- /dev/null +++ b/tests/unit/test_pipeline_scheduler.py @@ -0,0 +1,400 @@ +"""Unit tests for astrbot.core.pipeline.scheduler.PipelineScheduler.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, call, patch + +import pytest + +from astrbot.core.pipeline.scheduler import PipelineScheduler + + +@pytest.fixture +def mock_context(): + """Create a mock PipelineContext.""" + ctx = MagicMock() + ctx.astrbot_config = {"provider_settings": {"enable": True}} + ctx.plugin_manager = MagicMock() + ctx.plugin_manager.context = MagicMock(spec_set=[]) + return ctx + + +@pytest.fixture +def mock_stage_cls(): + """Create a mock stage class that can be instantiated.""" + cls = MagicMock() + instance = MagicMock() + instance.initialize = AsyncMock() + instance.process = AsyncMock() + cls.return_value = instance + cls.__name__ = "MockStage" + return cls + + +# --------------------------------------------------------------------------- +# Init +# --------------------------------------------------------------------------- + + +class TestPipelineSchedulerInit: + """Tests for PipelineScheduler.__init__().""" + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.STAGES_ORDER", ["MockStage"]) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + def test_init_calls_ensure_and_sets_context( + self, mock_ensure, mock_stages_order, mock_registered, mock_context, + ): + """Verify __init__ calls ensure_builtin_stages_registered and sets context.""" + mock_stage = MagicMock() + mock_stage.__name__ = "MockStage" + mock_registered.append(mock_stage) + + scheduler = PipelineScheduler(mock_context) + + mock_ensure.assert_called_once() + assert scheduler.ctx is mock_context + assert scheduler.stages == [] + + +# --------------------------------------------------------------------------- +# Initialize +# --------------------------------------------------------------------------- + + +class TestPipelineSchedulerInitialize: + """Tests for PipelineScheduler.initialize().""" + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_initialize_creates_stage_instances( + self, mock_ensure, mock_registered, mock_context, mock_stage_cls, + ): + """Verify initialize creates and initializes all registered stage instances.""" + mock_registered.append(mock_stage_cls) + scheduler = PipelineScheduler(mock_context) + + await scheduler.initialize() + + assert len(scheduler.stages) == 1 + assert scheduler.stages[0] is mock_stage_cls.return_value + mock_stage_cls.return_value.initialize.assert_awaited_once_with(mock_context) + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_initialize_multiple_stages( + self, mock_ensure, mock_registered, mock_context, + ): + """Verify multiple stages are initialized in order.""" + stage1_cls = MagicMock() + stage1_cls.__name__ = "Stage1" + stage1_instance = MagicMock() + stage1_instance.initialize = AsyncMock() + stage1_cls.return_value = stage1_instance + + stage2_cls = MagicMock() + stage2_cls.__name__ = "Stage2" + stage2_instance = MagicMock() + stage2_instance.initialize = AsyncMock() + stage2_cls.return_value = stage2_instance + + mock_registered.extend([stage1_cls, stage2_cls]) + scheduler = PipelineScheduler(mock_context) + + await scheduler.initialize() + + assert len(scheduler.stages) == 2 + stage1_instance.initialize.assert_awaited_once_with(mock_context) + stage2_instance.initialize.assert_awaited_once_with(mock_context) + + +# --------------------------------------------------------------------------- +# _process_stages +# --------------------------------------------------------------------------- + + +class TestPipelineSchedulerProcessStages: + """Tests for PipelineScheduler._process_stages().""" + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_non_generator_stage_executed( + self, mock_ensure, mock_registered, mock_context, mock_stage_cls, + ): + """Verify a non-generator stage is awaited.""" + mock_registered.append(mock_stage_cls) + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [mock_stage_cls.return_value] + event = MagicMock() + event.is_stopped.return_value = False + + await scheduler._process_stages(event) + + mock_stage_cls.return_value.process.assert_called_once_with(event) + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_generator_stage_yields_and_next_stage_runs( + self, mock_ensure, mock_registered, mock_context, + ): + """Verify a generator stage yields and the next stage runs.""" + async def gen_process(_event): + yield None + + stage1 = MagicMock() + stage1.process = gen_process + stage1.__class__.__name__ = "Stage1" + + stage2 = MagicMock() + stage2.process = AsyncMock() + stage2.__class__.__name__ = "Stage2" + + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [stage1, stage2] + event = MagicMock() + event.is_stopped.return_value = False + + await scheduler._process_stages(event) + + stage2.process.assert_called_once_with(event) + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_generator_stage_stops_propagation( + self, mock_ensure, mock_registered, mock_context, + ): + """Verify event.stop_event() breaks the pipeline in generator stage.""" + stage1_pass = [False] + + async def onion_process(_event): + stage1_pass[0] = True + yield None + + stage1 = MagicMock() + stage1.process = onion_process + stage1.__class__.__name__ = "Stage1" + + stage2 = MagicMock() + stage2.process = AsyncMock() + stage2.__class__.__name__ = "Stage2" + + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [stage1, stage2] + event = MagicMock() + event.is_stopped = lambda: True # always stopped + + await scheduler._process_stages(event) + + stage2.process.assert_not_called() + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_non_generator_stage_stops_propagation( + self, mock_ensure, mock_registered, mock_context, + ): + """Verify event.stop_event() breaks non-generator stage chain.""" + stage1 = MagicMock() + stage1.process = AsyncMock() + stage1.__class__.__name__ = "Stage1" + + stage2 = MagicMock() + stage2.process = AsyncMock() + stage2.__class__.__name__ = "Stage2" + + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [stage1, stage2] + event = MagicMock() + event.is_stopped = lambda: True # always stopped + + await scheduler._process_stages(event) + + stage1.process.assert_called_once_with(event) + stage2.process.assert_not_called() + + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_generator_with_onion_recursion( + self, mock_ensure, mock_registered, mock_context, + ): + """Verify generator stages recursively process subsequent stages (onion model).""" + call_order = [] + + async def onion_process(event): + call_order.append("before_yield") + yield None + call_order.append("after_yield") + + stage1 = MagicMock() + stage1.process = onion_process + stage1.__class__.__name__ = "Stage1" + + stage2 = MagicMock() + stage2.process = AsyncMock(side_effect=lambda e: call_order.append("stage2")) + stage2.__class__.__name__ = "Stage2" + + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [stage1, stage2] + event = MagicMock() + event.is_stopped.return_value = False + + await scheduler._process_stages(event) + + # Order: stage1 before_yield -> yield -> stage2 -> stage1 after_yield + assert call_order == ["before_yield", "stage2", "after_yield"] + + +# --------------------------------------------------------------------------- +# execute +# --------------------------------------------------------------------------- + + +class TestPipelineSchedulerExecute: + """Tests for PipelineScheduler.execute().""" + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_calls_process_stages( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify execute calls _process_stages and cleans up.""" + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock() + + event = MagicMock() + await scheduler.execute(event) + + scheduler._process_stages.assert_awaited_once_with(event) + mock_registry.register.assert_called_once_with(event) + mock_registry.unregister.assert_called_once_with(event) + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_webchat_event_sends_none( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify WebChatMessageEvent gets an extra None send.""" + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock() + + event = MagicMock(spec=["send", "is_stopped"]) + event.__class__.__name__ = "WebChatMessageEvent" + + with patch( + "astrbot.core.pipeline.scheduler.WebChatMessageEvent", + event.__class__, + ): + await scheduler.execute(event) + + event.send.assert_awaited_once_with(None) + mock_registry.unregister.assert_called_once_with(event) + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_wecom_event_sends_none( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify WecomAIBotMessageEvent gets an extra None send.""" + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock() + + event = MagicMock(spec=["send", "is_stopped"]) + event.__class__.__name__ = "WecomAIBotMessageEvent" + + with patch( + "astrbot.core.pipeline.scheduler.WecomAIBotMessageEvent", + event.__class__, + ): + await scheduler.execute(event) + + event.send.assert_awaited_once_with(None) + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_normal_event_no_extra_send( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify regular events do not get extra None send.""" + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock() + + event = MagicMock() + event.__class__.__name__ = "NormalEvent" + + await scheduler.execute(event) + + event.send.assert_not_called() + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_with_sdk_plugin_bridge( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify sdk_plugin_bridge.close_request_overlay_for_event is called.""" + mock_bridge = MagicMock() + mock_bridge.close_request_overlay_for_event = MagicMock() + mock_context.plugin_manager.context.sdk_plugin_bridge = mock_bridge + + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock() + + event = MagicMock() + await scheduler.execute(event) + + mock_bridge.close_request_overlay_for_event.assert_called_once_with(event) + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_without_sdk_plugin_bridge( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify no error when sdk_plugin_bridge is absent.""" + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock() + + event = MagicMock() + await scheduler.execute(event) + + # Should not raise + + @patch("astrbot.core.pipeline.scheduler.active_event_registry") + @patch("astrbot.core.pipeline.scheduler.registered_stages", []) + @patch("astrbot.core.pipeline.scheduler.ensure_builtin_stages_registered") + @pytest.mark.asyncio + async def test_execute_unregisters_on_error( + self, mock_ensure, mock_registered, mock_registry, mock_context, + ): + """Verify event is still unregistered when _process_stages raises.""" + scheduler = PipelineScheduler(mock_context) + scheduler.stages = [] + scheduler._process_stages = AsyncMock(side_effect=RuntimeError("fail")) + + event = MagicMock() + with pytest.raises(RuntimeError): + await scheduler.execute(event) + + mock_registry.unregister.assert_called_once_with(event) diff --git a/tests/unit/test_platform.py b/tests/unit/test_platform.py new file mode 100644 index 000000000..f38a5a6ed --- /dev/null +++ b/tests/unit/test_platform.py @@ -0,0 +1,312 @@ +"""Tests for astrbot.core.platform.platform — Platform ABC.""" + +import asyncio +from asyncio import Queue +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.platform.platform import Platform, PlatformError, PlatformStatus +from astrbot.core.platform.platform_metadata import PlatformMetadata + + +# --------------------------------------------------------------------------- +# Concrete subclass so we can test non-abstract behaviour of Platform. +# --------------------------------------------------------------------------- +class ConcretePlatform(Platform): + """Minimal concrete Platform used in every test that needs an instance.""" + + async def run(self) -> None: + pass + + def meta(self) -> PlatformMetadata: + return PlatformMetadata( + name="test_adapter", + description="A test adapter", + id="test_adapter_id", + adapter_display_name="Test Adapter", + ) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def config() -> dict: + return { + "key": "value", + "unified_webhook_mode": False, + } + + +@pytest.fixture +def event_queue() -> Queue: + return Queue() + + +@pytest.fixture +def platform(config: dict, event_queue: Queue) -> ConcretePlatform: + return ConcretePlatform(config, event_queue) + + +# =================================================================== +# Construction +# =================================================================== + +class TestConstruction: + """Platform.__init__ stores constructor arguments and sets defaults.""" + + def test_stores_config_and_queue(self, config: dict, event_queue: Queue): + p = ConcretePlatform(config, event_queue) + assert p.config is config + assert p._event_queue is event_queue + + def test_default_status_is_pending(self, platform: ConcretePlatform): + assert platform.status == PlatformStatus.PENDING + + def test_client_self_id_is_random_hex(self, platform: ConcretePlatform): + assert isinstance(platform.client_self_id, str) + assert len(platform.client_self_id) == 32 # uuid4 hex + + def test_errors_list_starts_empty(self, platform: ConcretePlatform): + assert platform.errors == [] + + def test_started_at_is_none_until_running(self, platform: ConcretePlatform): + assert platform._started_at is None + + +# =================================================================== +# Status property +# =================================================================== + +class TestStatus: + """Platform.status getter/setter and side-effects.""" + + def test_setter_changes_status(self, platform: ConcretePlatform): + platform.status = PlatformStatus.RUNNING + assert platform.status == PlatformStatus.RUNNING + + def test_setting_running_records_started_at(self, platform: ConcretePlatform): + platform.status = PlatformStatus.RUNNING + assert platform._started_at is not None + assert isinstance(platform._started_at, datetime) + + def test_setting_running_twice_does_not_overwrite_started_at( + self, platform: ConcretePlatform + ): + platform.status = PlatformStatus.RUNNING + t1 = platform._started_at + platform.status = PlatformStatus.ERROR + platform.status = PlatformStatus.RUNNING + assert platform._started_at == t1 + + +# =================================================================== +# Error paths +# =================================================================== + +class TestErrors: + """record_error, last_error, clear_errors.""" + + def test_record_error_appends_to_list(self, platform: ConcretePlatform): + platform.record_error("something went wrong", "traceback line 1") + assert len(platform.errors) == 1 + assert platform.errors[0].message == "something went wrong" + assert platform.errors[0].traceback == "traceback line 1" + + def test_record_error_sets_status_to_error(self, platform: ConcretePlatform): + platform.record_error("fail") + assert platform.status == PlatformStatus.ERROR + + def test_last_error_none_when_empty(self, platform: ConcretePlatform): + assert platform.last_error is None + + def test_last_error_returns_most_recent(self, platform: ConcretePlatform): + platform.record_error("first") + platform.record_error("second") + assert platform.last_error.message == "second" + + def test_clear_errors_empties_list(self, platform: ConcretePlatform): + platform.record_error("first") + platform.clear_errors() + assert platform.errors == [] + + def test_clear_errors_resets_status_from_error(self, platform: ConcretePlatform): + platform.record_error("first") + platform.clear_errors() + assert platform.status == PlatformStatus.RUNNING + + def test_clear_errors_is_noop_when_no_error(self, platform: ConcretePlatform): + platform.clear_errors() + assert platform.errors == [] + assert platform.status == PlatformStatus.PENDING + + +# =================================================================== +# unified_webhook +# =================================================================== + +class TestUnifiedWebhook: + """Platform.unified_webhook() logic.""" + + def test_disabled_by_default(self, platform: ConcretePlatform): + assert platform.unified_webhook() is False + + def test_enabled_when_both_config_present( + self, config: dict, event_queue: Queue + ): + config["unified_webhook_mode"] = True + config["webhook_uuid"] = "abc123" + p = ConcretePlatform(config, event_queue) + assert p.unified_webhook() is True + + def test_disabled_without_uuid(self, config: dict, event_queue: Queue): + config["unified_webhook_mode"] = True + # no webhook_uuid set + p = ConcretePlatform(config, event_queue) + assert p.unified_webhook() is False + + def test_disabled_when_mode_off(self, config: dict, event_queue: Queue): + config["unified_webhook_mode"] = False + config["webhook_uuid"] = "abc123" + p = ConcretePlatform(config, event_queue) + assert p.unified_webhook() is False + + +# =================================================================== +# get_stats +# =================================================================== + +class TestGetStats: + """Platform.get_stats() structure and content.""" + + def test_stats_contains_expected_keys(self, platform: ConcretePlatform): + stats = platform.get_stats() + expected_keys = { + "id", "type", "display_name", "status", "started_at", + "error_count", "last_error", "unified_webhook", "meta", + } + assert expected_keys.issubset(stats.keys()) + + def test_stats_values_without_errors(self, platform: ConcretePlatform): + stats = platform.get_stats() + assert stats["id"] == "test_adapter_id" + assert stats["type"] == "test_adapter" + assert stats["display_name"] == "Test Adapter" + assert stats["status"] == PlatformStatus.PENDING.value + assert stats["started_at"] is None + assert stats["error_count"] == 0 + assert stats["last_error"] is None + assert stats["unified_webhook"] is False + assert stats["meta"]["id"] == "test_adapter_id" + assert stats["meta"]["name"] == "test_adapter" + + def test_stats_reflects_recorded_errors(self, platform: ConcretePlatform): + platform.record_error("err1", "tb1") + platform.record_error("err2", "tb2") + stats = platform.get_stats() + assert stats["error_count"] == 2 + assert stats["last_error"]["message"] == "err2" + assert stats["last_error"]["traceback"] == "tb2" + + +# =================================================================== +# Instance methods +# =================================================================== + +class TestMethods: + """terminate, get_client, commit_event, webhook_callback, send_by_session.""" + + @pytest.mark.asyncio + async def test_terminate_sets_stopped(self, platform: ConcretePlatform): + await platform.terminate() + assert platform.status == PlatformStatus.STOPPED + + def test_get_client_returns_none(self, platform: ConcretePlatform): + assert platform.get_client() is None + + @pytest.mark.asyncio + async def test_commit_event_puts_into_queue( + self, platform: ConcretePlatform, event_queue: Queue + ): + mock_event = MagicMock() + platform.commit_event(mock_event) + assert event_queue.qsize() == 1 + assert await event_queue.get() is mock_event + + @pytest.mark.asyncio + async def test_webhook_callback_raises_not_implemented( + self, platform: ConcretePlatform + ): + with pytest.raises(NotImplementedError) as exc: + await platform.webhook_callback(None) + assert "未实现统一 Webhook 模式" in str(exc.value) + + @pytest.mark.asyncio + async def test_send_by_session_calls_metric_upload( + self, platform: ConcretePlatform + ): + mock_session = MagicMock() + mock_chain = MagicMock() + with patch( + "astrbot.core.platform.platform.Metric.upload", + new_callable=AsyncMock, + ) as mock_upload: + await platform.send_by_session(mock_session, mock_chain) + mock_upload.assert_called_once_with( + msg_event_tick=1, adapter_name="test_adapter" + ) + + +# =================================================================== +# Abstract-method detection (cannot instantiate ABC directly) +# =================================================================== + +class TestAbstractDetection: + """Verify the ABC enforces that run() and meta() are implemented.""" + + def test_cannot_instantiate_platform_directly(self): + with pytest.raises(TypeError): + Platform({"k": "v"}, Queue()) # type: ignore[abstract] + + def test_missing_run_raises_type_error(self): + class MissingRun(Platform): + def meta(self) -> PlatformMetadata: + return PlatformMetadata(name="x", description="x", id="x") + + with pytest.raises(TypeError): + MissingRun({}, Queue()) # type: ignore[abstract] + + def test_missing_meta_raises_type_error(self): + class MissingMeta(Platform): + async def run(self) -> None: + pass + + with pytest.raises(TypeError): + MissingMeta({}, Queue()) # type: ignore[abstract] + + +# =================================================================== +# PlatformError dataclass +# =================================================================== + +class TestPlatformErrorDataclass: + """PlatformError construction and defaults.""" + + def test_required_message(self): + err = PlatformError(message="oops") + assert err.message == "oops" + assert err.traceback is None + assert isinstance(err.timestamp, datetime) + + def test_with_traceback(self): + err = PlatformError(message="oops", traceback="tb content") + assert err.traceback == "tb content" + + def test_default_timestamp_is_nowish(self): + before = datetime.now() + err = PlatformError(message="oops") + after = datetime.now() + assert before <= err.timestamp <= after diff --git a/tests/unit/test_provider_entities.py b/tests/unit/test_provider_entities.py new file mode 100644 index 000000000..469385e39 --- /dev/null +++ b/tests/unit/test_provider_entities.py @@ -0,0 +1,518 @@ +"""Tests for astrbot.core.provider.entities. + +Covers ProviderRequest, LLMResponse, TokenUsage, ToolCallsResult, +ProviderMeta, ProviderMetaData, and RerankResult construction and edge cases. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.agent.message import ( + AssistantMessageSegment, + ContentPart, + ToolCall, + ToolCallMessageSegment, +) +from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.provider.entities import ( + LLMResponse, + ProviderMeta, + ProviderMetaData, + ProviderRequest, + ProviderType, + RerankResult, + TokenUsage, + ToolCallsResult, +) + + +# ========================================================================= +# ProviderMeta / ProviderMetaData +# ========================================================================= + + +class TestProviderMeta: + def test_basic_construction(self): + meta = ProviderMeta(id="p1", model="gpt-4", type="openai") + assert meta.id == "p1" + assert meta.model == "gpt-4" + assert meta.type == "openai" + assert meta.provider_type == ProviderType.CHAT_COMPLETION + + def test_construction_with_provider_type(self): + meta = ProviderMeta( + id="emb1", + model="text-embedding-3", + type="openai_embedding", + provider_type=ProviderType.EMBEDDING, + ) + assert meta.provider_type == ProviderType.EMBEDDING + + +class TestProviderMetaData: + def test_basic_construction(self): + pmd = ProviderMetaData( + id="p1", model=None, type="openai", desc="OpenAI provider" + ) + assert pmd.id == "p1" + assert pmd.desc == "OpenAI provider" + assert pmd.cls_type is None + assert pmd.default_config_tmpl is None + assert pmd.provider_display_name is None + + def test_construction_with_all_fields(self): + fake_cls = type("FakeProvider", (), {}) + pmd = ProviderMetaData( + id="p2", + model="gpt-4o", + type="openai", + desc="desc", + provider_type=ProviderType.CHAT_COMPLETION, + cls_type=fake_cls, + default_config_tmpl={"key": "val"}, + provider_display_name="OpenAI Official", + ) + assert pmd.cls_type is fake_cls + assert pmd.default_config_tmpl == {"key": "val"} + assert pmd.provider_display_name == "OpenAI Official" + + +# ========================================================================= +# TokenUsage +# ========================================================================= + + +class TestTokenUsage: + def test_defaults(self): + tu = TokenUsage() + assert tu.input_other == 0 + assert tu.input_cached == 0 + assert tu.output == 0 + assert tu.total == 0 + assert tu.input == 0 + + def test_properties(self): + tu = TokenUsage(input_other=10, input_cached=5, output=20) + assert tu.total == 35 + assert tu.input == 15 + + def test_addition(self): + a = TokenUsage(input_other=5, input_cached=2, output=10) + b = TokenUsage(input_other=3, input_cached=1, output=4) + result = a + b + assert result.input_other == 8 + assert result.input_cached == 3 + assert result.output == 14 + + def test_subtraction(self): + a = TokenUsage(input_other=10, input_cached=5, output=20) + b = TokenUsage(input_other=3, input_cached=2, output=5) + result = a - b + assert result.input_other == 7 + assert result.input_cached == 3 + assert result.output == 15 + + def test_addition_preserves_immutability(self): + a = TokenUsage(input_other=1, output=2) + b = TokenUsage(input_other=3, output=4) + c = a + b + assert a.input_other == 1 + assert b.input_other == 3 + assert c.input_other == 4 + + +# ========================================================================= +# ToolCallsResult +# ========================================================================= + + +class TestToolCallsResult: + def test_construction(self): + info = MagicMock(spec=AssistantMessageSegment) + info.model_dump.return_value = {"role": "assistant", "content": "thinking"} + result_seg = MagicMock(spec=ToolCallMessageSegment) + result_seg.model_dump.return_value = {"role": "tool", "content": "result"} + + tcr = ToolCallsResult( + tool_calls_info=info, + tool_calls_result=[result_seg], + ) + assert tcr.tool_calls_info is info + assert len(tcr.tool_calls_result) == 1 + + def test_to_openai_messages(self): + info = MagicMock(spec=AssistantMessageSegment) + info.model_dump.return_value = {"role": "assistant"} + r1 = MagicMock(spec=ToolCallMessageSegment) + r1.model_dump.return_value = {"role": "tool", "name": "get_weather"} + r2 = MagicMock(spec=ToolCallMessageSegment) + r2.model_dump.return_value = {"role": "tool", "name": "search"} + + tcr = ToolCallsResult(tool_calls_info=info, tool_calls_result=[r1, r2]) + msgs = tcr.to_openai_messages() + assert len(msgs) == 3 + assert msgs[0] == {"role": "assistant"} + assert msgs[1] == {"role": "tool", "name": "get_weather"} + assert msgs[2] == {"role": "tool", "name": "search"} + + def test_to_openai_messages_model(self): + info = MagicMock(spec=AssistantMessageSegment) + r1 = MagicMock(spec=ToolCallMessageSegment) + + tcr = ToolCallsResult(tool_calls_info=info, tool_calls_result=[r1]) + models = tcr.to_openai_messages_model() + assert len(models) == 2 + assert models[0] is info + assert models[1] is r1 + + def test_to_openai_messages_empty_result(self): + info = MagicMock(spec=AssistantMessageSegment) + info.model_dump.return_value = {"role": "assistant"} + tcr = ToolCallsResult(tool_calls_info=info, tool_calls_result=[]) + msgs = tcr.to_openai_messages() + assert len(msgs) == 1 + assert msgs[0] == {"role": "assistant"} + + +# ========================================================================= +# ProviderRequest +# ========================================================================= + + +class TestProviderRequest: + def test_default_construction(self): + req = ProviderRequest() + assert req.prompt is None + assert req.session_id == "" + assert req.image_urls == [] + assert req.audio_urls == [] + assert req.contexts == [] + assert req.func_tool is None + assert req.system_prompt is None + assert req.conversation is None + assert req.tool_calls_result is None + assert req.model is None + + def test_construction_with_values(self): + req = ProviderRequest( + prompt="Hello", + session_id="sess-1", + image_urls=["http://example.com/img.png"], + system_prompt="You are a bot", + model="gpt-4", + ) + assert req.prompt == "Hello" + assert req.session_id == "sess-1" + assert req.image_urls == ["http://example.com/img.png"] + assert req.system_prompt == "You are a bot" + assert req.model == "gpt-4" + + def test_repr_without_context(self): + req = ProviderRequest(prompt="hi", session_id="s1") + text = repr(req) + assert "prompt=hi" in text + assert "session_id=s1" in text + assert "image_count=0" in text + + def test_repr_with_conversation(self): + conversation = MagicMock() + conversation.cid = "conv-abc" + req = ProviderRequest(prompt="test", conversation=conversation) + text = repr(req) + assert "conversation_id=conv-abc" in text + + def test_append_tool_calls_result_none_to_single(self): + req = ProviderRequest() + tcr = MagicMock(spec=ToolCallsResult) + req.append_tool_calls_result(tcr) + assert isinstance(req.tool_calls_result, list) + assert len(req.tool_calls_result) == 1 + assert req.tool_calls_result[0] is tcr + + def test_append_tool_calls_result_single_to_list(self): + tcr1 = MagicMock(spec=ToolCallsResult) + req = ProviderRequest(tool_calls_result=tcr1) + tcr2 = MagicMock(spec=ToolCallsResult) + req.append_tool_calls_result(tcr2) + assert isinstance(req.tool_calls_result, list) + assert len(req.tool_calls_result) == 2 + assert req.tool_calls_result[0] is tcr1 + assert req.tool_calls_result[1] is tcr2 + + def test_append_tool_calls_result_list(self): + tcr1 = MagicMock(spec=ToolCallsResult) + tcr2 = MagicMock(spec=ToolCallsResult) + req = ProviderRequest(tool_calls_result=[tcr1]) + tcr3 = MagicMock(spec=ToolCallsResult) + req.append_tool_calls_result(tcr3) + assert len(req.tool_calls_result) == 2 + + def test_print_friendly_context_no_contexts(self): + req = ProviderRequest(prompt="hello", image_urls=["a.png"], audio_urls=["b.wav"]) + result = req._print_friendly_context() + assert "prompt: hello" in result + assert "image_count: 1" in result + assert "audio_count: 1" in result + + def test_print_friendly_context_with_text_contexts(self): + req = ProviderRequest(contexts=[{"role": "user", "content": "hello"}]) + result = req._print_friendly_context() + assert "user: hello" in result + + def test_print_friendly_context_filters_checkpoints(self): + req = ProviderRequest( + contexts=[ + {"role": "user", "content": "hi", "checkpoint": True}, + {"role": "assistant", "content": "hello"}, + ] + ) + with patch( + "astrbot.core.provider.entities.is_checkpoint_message", + side_effect=lambda c: c.get("checkpoint", False), + ): + result = req._print_friendly_context() + assert "user: hi" not in result + assert "assistant: hello" in result + + def test_print_friendly_context_multimodal(self): + req = ProviderRequest( + contexts=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "x.jpg"}}, + {"type": "image_url", "image_url": {"url": "y.jpg"}}, + {"type": "audio_url", "audio_url": {"url": "z.wav"}}, + ], + } + ] + ) + result = req._print_friendly_context() + assert "user: describe this[+2 images][+1 audios]" in result + + def test_assemble_context_simple_text(self): + """When there's only a plain text prompt and no extra content, it returns a simple str content.""" + req = ProviderRequest(prompt="hello world") + import asyncio + + ctx = asyncio.run(req.assemble_context()) + assert ctx == {"role": "user", "content": "hello world"} + + def test_assemble_context_empty_prompt_with_images_adds_placeholder(self): + req = ProviderRequest(prompt="", image_urls=[]) + import asyncio + + with patch.object(req, "_encode_image_bs64", return_value="data:image/jpeg;base64,abc"): + req.image_urls = ["http://example.com/img.png"] + ctx = asyncio.run(req.assemble_context()) + assert ctx["role"] == "user" + # Should include "[图片]" placeholder + content = ctx["content"] + assert isinstance(content, list) + assert any(b.get("text") == "[图片]" for b in content) + + def test_assemble_context_empty_prompt_with_audio_adds_placeholder(self): + req = ProviderRequest(prompt="", audio_urls=["http://example.com/a.wav"]) + import asyncio + + with ( + patch.object(req, "_encode_audio_bs64", return_value="data:audio/wav;base64,xyz"), + patch("astrbot.core.provider.entities.download_file", AsyncMock()), + ): + ctx = asyncio.run(req.assemble_context()) + assert ctx["role"] == "user" + content = ctx["content"] + assert isinstance(content, list) + assert any(b.get("text") == "[音频]" for b in content) + + def test_assemble_context_with_extra_user_content(self): + extra_part = MagicMock(spec=ContentPart) + extra_part.model_dump.return_value = {"type": "text", "text": "extra instruction"} + req = ProviderRequest( + prompt="translate this", + extra_user_content_parts=[extra_part], + ) + import asyncio + + ctx = asyncio.run(req.assemble_context()) + assert ctx["role"] == "user" + content = ctx["content"] + assert isinstance(content, list) + texts = [b["text"] for b in content if b.get("type") == "text"] + assert "translate this" in texts + assert "extra instruction" in texts + + def test_encode_image_bs64_base64_prefix(self): + req = ProviderRequest() + import asyncio + + result = asyncio.run(req._encode_image_bs64("base64://rawdata")) + assert result == "data:image/jpeg;base64,rawdata" + + def test_encode_audio_bs64_base64_prefix(self): + req = ProviderRequest() + import asyncio + + result = asyncio.run(req._encode_audio_bs64("base64://rawdata")) + assert result == "data:audio/wav;base64,rawdata" + + def test_str_equals_repr(self): + req = ProviderRequest(prompt="test", session_id="s1") + assert str(req) == repr(req) + + +# ========================================================================= +# LLMResponse +# ========================================================================= + + +class TestLLMResponse: + def test_default_role_assistant(self): + resp = LLMResponse(role="assistant") + assert resp.role == "assistant" + assert resp.completion_text is None or resp.completion_text == "" + assert resp.tools_call_args == [] + assert resp.tools_call_name == [] + assert resp.tools_call_ids == [] + assert resp.tools_call_extra_content == {} + assert resp.reasoning_content is None + assert resp.raw_completion is None + assert resp.is_chunk is False + assert resp.id is None + assert resp.usage is None + + def test_construction_with_completion_text(self): + resp = LLMResponse(role="assistant", completion_text="Hello world") + assert resp.completion_text == "Hello world" + + def test_construction_with_result_chain(self): + chain = MessageChain() + chain.message("Hello from chain") + resp = LLMResponse(role="assistant", result_chain=chain) + assert resp.completion_text == "Hello from chain" + + def test_completion_text_setter_with_result_chain(self): + chain = MessageChain() + chain.message("Old text") + resp = LLMResponse(role="assistant", result_chain=chain) + assert resp.completion_text == "Old text" + resp.completion_text = "New text" + # The setter inserts a Plain component at the start after removing old ones + assert "New text" in resp.completion_text + + def test_completion_text_setter_without_result_chain(self): + resp = LLMResponse(role="assistant") + resp.completion_text = "direct text" + assert resp.completion_text == "direct text" + + def test_tool_calls_defaults_to_empty(self): + resp = LLMResponse(role="assistant") + # They should be empty lists, not None + assert resp.tools_call_args == [] + assert resp.tools_call_name == [] + assert resp.tools_call_ids == [] + assert resp.tools_call_extra_content == {} + + def test_construction_with_tool_calls(self): + resp = LLMResponse( + role="assistant", + tools_call_args=[{"location": "NYC"}], + tools_call_name=["get_weather"], + tools_call_ids=["call_123"], + tools_call_extra_content={"call_123": {"source": "web"}}, + ) + assert resp.tools_call_args == [{"location": "NYC"}] + assert resp.tools_call_name == ["get_weather"] + assert resp.tools_call_ids == ["call_123"] + assert resp.tools_call_extra_content == {"call_123": {"source": "web"}} + + def test_to_openai_tool_calls(self): + resp = LLMResponse( + role="assistant", + tools_call_args=[{"q": "weather"}, {"q": "news"}], + tools_call_name=["search", "search"], + tools_call_ids=["c1", "c2"], + tools_call_extra_content={"c1": {"priority": 1}}, + ) + calls = resp.to_openai_tool_calls() + assert len(calls) == 2 + assert calls[0]["id"] == "c1" + assert calls[0]["function"]["name"] == "search" + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "weather"} + assert calls[0]["extra_content"] == {"priority": 1} + assert calls[1]["id"] == "c2" + assert "extra_content" not in calls[1] + + def test_to_openai_to_calls_model(self): + resp = LLMResponse( + role="assistant", + tools_call_args=[{"x": 1}], + tools_call_name=["foo"], + tools_call_ids=["cid1"], + tools_call_extra_content={"cid1": {"meta": "data"}}, + ) + calls = resp.to_openai_to_calls_model() + assert len(calls) == 1 + assert isinstance(calls[0], ToolCall) + assert calls[0].id == "cid1" + assert calls[0].function.name == "foo" + assert calls[0].extra_content == {"meta": "data"} + + def test_to_openai_tool_calls_empty(self): + resp = LLMResponse(role="assistant") + calls = resp.to_openai_tool_calls() + assert calls == [] + + def test_to_openai_to_calls_model_empty(self): + resp = LLMResponse(role="assistant") + calls = resp.to_openai_to_calls_model() + assert calls == [] + + def test_construction_with_reasoning(self): + resp = LLMResponse( + role="assistant", + completion_text="final answer", + reasoning_content="thinking step by step", + reasoning_signature="sig_abc", + ) + assert resp.reasoning_content == "thinking step by step" + assert resp.reasoning_signature == "sig_abc" + + def test_construction_with_raw_completion(self): + raw = MagicMock() + resp = LLMResponse(role="assistant", raw_completion=raw) + assert resp.raw_completion is raw + + def test_construction_with_usage(self): + usage = TokenUsage(input_other=10, output=20) + resp = LLMResponse(role="assistant", usage=usage) + assert resp.usage is usage + + def test_construction_with_chunk_and_id(self): + resp = LLMResponse(role="assistant", is_chunk=True, id="chunk_1") + assert resp.is_chunk is True + assert resp.id == "chunk_1" + + +# ========================================================================= +# RerankResult +# ========================================================================= + + +class TestRerankResult: + def test_construction(self): + rr = RerankResult(index=0, relevance_score=0.95) + assert rr.index == 0 + assert rr.relevance_score == 0.95 + + def test_negative_score(self): + rr = RerankResult(index=5, relevance_score=-0.1) + assert rr.relevance_score == -0.1 + + def test_zero_values(self): + rr = RerankResult(index=0, relevance_score=0.0) + assert rr.index == 0 + assert rr.relevance_score == 0.0 diff --git a/tests/unit/test_provider_manager.py b/tests/unit/test_provider_manager.py new file mode 100644 index 000000000..e5bf58eb6 --- /dev/null +++ b/tests/unit/test_provider_manager.py @@ -0,0 +1,408 @@ +"""Tests for astrbot.core.provider.manager. + +Covers ProviderManager __init__, callbacks/hooks, get_provider_by_id, +get_using_provider, _resolve_env_key_list, get_provider_config_by_id, +and related helper methods. +""" + +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.provider.entities import ProviderType +from astrbot.core.provider.manager import ProviderManager +from astrbot.core.provider.provider import ( + EmbeddingProvider, + Provider, + RerankProvider, + STTProvider, + TTSProvider, +) + + +# ========================================================================= +# Fixtures +# ========================================================================= + + +@pytest.fixture +def mock_acm(): + """Create a mock AstrBotConfigManager with valid nested config structure.""" + acm = MagicMock() + acm.confs = { + "default": { + "provider": [], + "provider_sources": [], + "provider_settings": {"default_provider_id": ""}, + "provider_stt_settings": {}, + "provider_tts_settings": {}, + } + } + acm.default_conf = acm.confs["default"] + acm.get_conf.return_value = acm.confs["default"] + return acm + + +@pytest.fixture +def mock_db(): + return MagicMock() + + +@pytest.fixture +def mock_persona_mgr(): + pm = MagicMock() + pm.default_persona = "default" + pm.persona_v3_config = [] + pm.personas_v3 = [] + pm.selected_default_persona_v3 = None + return pm + + +@pytest.fixture +def manager(mock_acm, mock_db, mock_persona_mgr): + with patch("astrbot.core.provider.manager.llm_tools") as mock_llm_tools: + mgr = ProviderManager( + acm=mock_acm, + db_helper=mock_db, + persona_mgr=mock_persona_mgr, + ) + yield mgr + + +# ========================================================================= +# __init__ +# ========================================================================= + + +class TestProviderManagerInit: + def test_construction(self, manager): + assert manager.reload_lock is not None + assert manager.resource_lock is not None + assert manager.providers_config == [] + assert manager.provider_sources_config == [] + assert manager.provider_settings == {"default_provider_id": ""} + assert manager.provider_insts == [] + assert manager.stt_provider_insts == [] + assert manager.tts_provider_insts == [] + assert manager.embedding_provider_insts == [] + assert manager.rerank_provider_insts == [] + assert manager.inst_map == {} + assert manager.curr_provider_inst is None + assert manager._provider_change_callback is None + assert manager._provider_change_hooks == [] + assert manager._mcp_init_task is None + + def test_default_persona_name_from_mgr(self, manager): + assert manager.default_persona_name == "default" + + def test_persona_configs_property(self, manager): + assert manager.persona_configs == [] + + def test_personas_property(self, manager): + assert manager.personas == [] + + def test_selected_default_persona_property(self, manager): + assert manager.selected_default_persona is None + + +# ========================================================================= +# Callbacks / Hooks +# ========================================================================= + + +class TestProviderManagerCallbacks: + def test_set_provider_change_callback(self, manager): + cb = MagicMock() + manager.set_provider_change_callback(cb) + assert manager._provider_change_callback is cb + + def test_set_provider_change_callback_none(self, manager): + manager.set_provider_change_callback(None) + assert manager._provider_change_callback is None + + def test_register_provider_change_hook(self, manager): + hook = MagicMock() + manager.register_provider_change_hook(hook) + assert hook in manager._provider_change_hooks + + def test_register_provider_change_hook_duplicate(self, manager): + hook = MagicMock() + manager.register_provider_change_hook(hook) + manager.register_provider_change_hook(hook) + assert len(manager._provider_change_hooks) == 1 + + def test_unregister_provider_change_hook(self, manager): + hook = MagicMock() + manager.register_provider_change_hook(hook) + manager.unregister_provider_change_hook(hook) + assert hook not in manager._provider_change_hooks + + def test_unregister_provider_change_hook_not_registered(self, manager): + hook = MagicMock() + # Should not raise + manager.unregister_provider_change_hook(hook) + + def test_notify_provider_changed_calls_callback(self, manager): + cb = MagicMock() + manager.set_provider_change_callback(cb) + manager._notify_provider_changed("p1", ProviderType.CHAT_COMPLETION, "umo_1") + cb.assert_called_once_with("p1", ProviderType.CHAT_COMPLETION, "umo_1") + + def test_notify_provider_changed_swallows_callback_error(self, manager): + cb = MagicMock(side_effect=ValueError("oops")) + manager.set_provider_change_callback(cb) + # Should not raise + manager._notify_provider_changed("p1", ProviderType.CHAT_COMPLETION, None) + + def test_notify_provider_changed_calls_hooks(self, manager): + hook1 = MagicMock() + hook2 = MagicMock() + manager.register_provider_change_hook(hook1) + manager.register_provider_change_hook(hook2) + manager._notify_provider_changed("p1", ProviderType.SPEECH_TO_TEXT, None) + hook1.assert_called_once_with("p1", ProviderType.SPEECH_TO_TEXT, None) + hook2.assert_called_once_with("p1", ProviderType.SPEECH_TO_TEXT, None) + + def test_notify_provider_changed_skips_callback_in_hooks(self, manager): + """When the same callable is both callback and hook, it should only be invoked once.""" + fn = MagicMock() + manager.set_provider_change_callback(fn) + manager.register_provider_change_hook(fn) + manager._notify_provider_changed("p1", ProviderType.CHAT_COMPLETION, None) + assert fn.call_count == 1 + + def test_notify_provider_changed_swallows_hook_error(self, manager): + hook = MagicMock(side_effect=RuntimeError("hook failed")) + manager.register_provider_change_hook(hook) + # Should not raise + manager._notify_provider_changed("p1", ProviderType.CHAT_COMPLETION, None) + + +# ========================================================================= +# get_provider_by_id / get_using_provider +# ========================================================================= + + +class TestProviderManagerLookups: + def test_get_provider_by_id_not_found(self, manager): + result = manager.get_provider_by_id("nonexistent") + assert result is None + + def test_get_provider_by_id_found(self, manager): + fake_inst = MagicMock(spec=Provider) + manager.inst_map["p1"] = fake_inst + result = manager.get_provider_by_id("p1") + assert result is fake_inst + + def test_get_using_provider_chat_completion_default(self, manager, mock_acm): + fake_provider = MagicMock(spec=Provider) + manager.provider_insts = [fake_provider] + mock_acm.get_conf.return_value = { + "provider_settings": {"default_provider_id": None} + } + result = manager.get_using_provider(ProviderType.CHAT_COMPLETION) + assert result is fake_provider + + def test_get_using_provider_chat_completion_by_id(self, manager, mock_acm): + fake_provider = MagicMock(spec=Provider) + manager.inst_map["default_prov"] = fake_provider + mock_acm.get_conf.return_value = { + "provider_settings": {"default_provider_id": "default_prov"} + } + result = manager.get_using_provider(ProviderType.CHAT_COMPLETION) + assert result is fake_provider + + def test_get_using_provider_chat_completion_no_instances(self, manager, mock_acm): + mock_acm.get_conf.return_value = { + "provider_settings": {"default_provider_id": None} + } + result = manager.get_using_provider(ProviderType.CHAT_COMPLETION) + assert result is None + + def test_get_using_provider_stt_disabled_returns_none(self, manager, mock_acm): + mock_acm.get_conf.return_value = { + "provider_stt_settings": {"enable": False} + } + result = manager.get_using_provider(ProviderType.SPEECH_TO_TEXT) + assert result is None + + def test_get_using_provider_tts_disabled_returns_none(self, manager, mock_acm): + mock_acm.get_conf.return_value = { + "provider_tts_settings": {"enable": False} + } + result = manager.get_using_provider(ProviderType.TEXT_TO_SPEECH) + assert result is None + + def test_get_using_provider_unknown_type(self, manager, mock_acm): + with pytest.raises(ValueError, match="Unknown provider type"): + manager.get_using_provider(ProviderType.EMBEDDING) + + def test_get_using_provider_with_umo(self, manager, mock_acm): + fake_provider = MagicMock(spec=Provider) + manager.inst_map["umo_prov"] = fake_provider + mock_acm.get_conf.return_value = { + "provider_settings": {"default_provider_id": None} + } + with patch("astrbot.core.provider.manager.sp") as mock_sp: + mock_sp.get.return_value = "umo_prov" + result = manager.get_using_provider( + ProviderType.CHAT_COMPLETION, umo="session_1" + ) + assert result is fake_provider + mock_sp.get.assert_called_once() + + +# ========================================================================= +# _resolve_env_key_list +# ========================================================================= + + +class TestResolveEnvKeyList: + def test_no_env_vars(self, manager): + config = {"key": ["sk-abc", "sk-def"]} + result = manager._resolve_env_key_list(config) + assert result["key"] == ["sk-abc", "sk-def"] + + def test_env_var_resolved(self, manager): + os.environ["MY_API_KEY"] = "sk-from-env" + config = {"key": ["$MY_API_KEY"], "id": "prov1"} + result = manager._resolve_env_key_list(config) + assert result["key"] == ["sk-from-env"] + os.environ.pop("MY_API_KEY", None) + + def test_env_var_braces(self, manager): + os.environ["SECRET"] = "very_secret" + config = {"key": ["${SECRET}"], "id": "prov1"} + result = manager._resolve_env_key_list(config) + assert result["key"] == ["very_secret"] + os.environ.pop("SECRET", None) + + def test_env_var_not_set_logs_warning(self, manager): + # Ensure env var does not exist + os.environ.pop("UNSET_VAR", None) + config = {"key": ["$UNSET_VAR"], "id": "prov1"} + with patch("astrbot.core.provider.manager.logger") as mock_logger: + result = manager._resolve_env_key_list(config) + assert result["key"] == [""] + mock_logger.warning.assert_called_once() + + def test_env_var_mixed_list(self, manager): + os.environ["K1"] = "val1" + config = {"key": ["$K1", "static_key"], "id": "prov1"} + result = manager._resolve_env_key_list(config) + assert result["key"] == ["val1", "static_key"] + os.environ.pop("K1", None) + + def test_non_list_key_passthrough(self, manager): + config = {"key": "not_a_list"} + result = manager._resolve_env_key_list(config) + assert result["key"] == "not_a_list" + + def test_empty_env_var_name(self, manager): + config = {"key": ["$"], "id": "prov1"} + result = manager._resolve_env_key_list(config) + assert result["key"] == [""] + + def test_missing_key_field(self, manager): + config = {"id": "prov1"} + result = manager._resolve_env_key_list(config) + assert result == config + + +# ========================================================================= +# get_provider_config_by_id +# ========================================================================= + + +class TestGetProviderConfigById: + def test_found(self, manager): + manager.providers_config = [ + {"id": "p1", "type": "openai"}, + {"id": "p2", "type": "anthropic"}, + ] + result = manager.get_provider_config_by_id("p1") + assert result == {"id": "p1", "type": "openai"} + + def test_not_found(self, manager): + manager.providers_config = [{"id": "p1", "type": "openai"}] + result = manager.get_provider_config_by_id("nonexistent") + assert result is None + + def test_deep_copy_returned(self, manager): + manager.providers_config = [{"id": "p1", "type": "openai", "key": ["secret"]}] + result = manager.get_provider_config_by_id("p1") + # Mutating the result should not affect the source + result["key"].append("new_key") + assert manager.providers_config[0]["key"] == ["secret"] + + def test_merged_flag(self, manager): + manager.providers_config = [ + {"id": "p1", "type": "openai", "provider_source_id": "src1"} + ] + manager.provider_sources_config = [ + {"id": "src1", "base_url": "https://api.openai.com"} + ] + with patch.object( + manager, + "get_merged_provider_config", + return_value={"id": "p1", "type": "openai", "base_url": "https://api.openai.com"}, + ): + result = manager.get_provider_config_by_id("p1", merged=True) + assert result["base_url"] == "https://api.openai.com" + + def test_empty_configs(self, manager): + result = manager.get_provider_config_by_id("p1") + assert result is None + + +# ========================================================================= +# _get_all_provider_instances / _clear_loaded_instances / get_insts +# ========================================================================= + + +class TestProviderManagerInstances: + def test_get_insts(self, manager): + fp = MagicMock(spec=Provider) + manager.provider_insts = [fp] + assert manager.get_insts() == [fp] + + def test_get_all_provider_instances_deduplicates(self, manager): + fp = MagicMock(spec=Provider) + manager.provider_insts = [fp] + manager.inst_map = {"p1": fp} + all_insts = manager._get_all_provider_instances() + # fp appears in both lists but should only be returned once + assert len(all_insts) == 1 + assert all_insts[0] is fp + + def test_get_all_provider_instances_returns_all_types(self, manager): + fp = MagicMock(spec=Provider) + stt = MagicMock(spec=STTProvider) + tts = MagicMock(spec=TTSProvider) + emb = MagicMock(spec=EmbeddingProvider) + rerank = MagicMock(spec=RerankProvider) + manager.provider_insts = [fp] + manager.stt_provider_insts = [stt] + manager.tts_provider_insts = [tts] + manager.embedding_provider_insts = [emb] + manager.rerank_provider_insts = [rerank] + all_insts = manager._get_all_provider_instances() + assert len(all_insts) == 5 + + def test_clear_loaded_instances(self, manager): + manager.provider_insts = [MagicMock(spec=Provider)] + manager.stt_provider_insts = [MagicMock(spec=STTProvider)] + manager.inst_map = {"p1": MagicMock()} + manager.curr_provider_inst = MagicMock() + + manager._clear_loaded_instances() + assert manager.provider_insts == [] + assert manager.stt_provider_insts == [] + assert manager.tts_provider_insts == [] + assert manager.embedding_provider_insts == [] + assert manager.rerank_provider_insts == [] + assert manager.inst_map == {} + assert manager.curr_provider_inst is None + assert manager.curr_stt_provider_inst is None + assert manager.curr_tts_provider_inst is None diff --git a/tests/unit/test_provider_provider.py b/tests/unit/test_provider_provider.py new file mode 100644 index 000000000..3ce0ac3b8 --- /dev/null +++ b/tests/unit/test_provider_provider.py @@ -0,0 +1,473 @@ +"""Tests for astrbot.core.provider.provider. + +Covers AbstractProvider, Provider, STTProvider, TTSProvider, +EmbeddingProvider, and RerankProvider abstract/concrete methods. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.agent.message import Message +from astrbot.core.provider.entities import ( + LLMResponse, + ProviderMeta, + ProviderMetaData, + ProviderType, +) +from astrbot.core.provider.provider import ( + AbstractProvider, + EmbeddingProvider, + Provider, + RerankProvider, + STTProvider, + TTSProvider, +) + + +# ========================================================================= +# AbstractProvider +# ========================================================================= + + +class TestAbstractProvider: + def test_construction(self): + ap = AbstractProvider(provider_config={"type": "test"}) + assert ap.model_name == "" + assert ap.provider_config == {"type": "test"} + + def test_set_and_get_model(self): + ap = AbstractProvider(provider_config={"type": "test"}) + assert ap.get_model() == "" + ap.set_model("gpt-4") + assert ap.get_model() == "gpt-4" + + def test_set_model_empty_string(self): + ap = AbstractProvider(provider_config={"type": "test"}) + ap.set_model("gpt-4") + ap.set_model("") + assert ap.get_model() == "" + + def test_meta_returns_provider_meta(self): + pmd = ProviderMetaData( + id="default", + model=None, + type="test_type", + provider_type=ProviderType.CHAT_COMPLETION, + ) + with patch( + "astrbot.core.provider.provider.provider_cls_map", + {"test_type": pmd}, + ): + ap = AbstractProvider(provider_config={"type": "test_type", "id": "myid"}) + meta = ap.meta() + assert isinstance(meta, ProviderMeta) + assert meta.id == "myid" + assert meta.type == "test_type" + assert meta.provider_type == ProviderType.CHAT_COMPLETION + + def test_meta_raises_on_unregistered_type(self): + ap = AbstractProvider(provider_config={"type": "nonexistent"}) + with pytest.raises(ValueError, match="not registered"): + ap.meta() + + def test_meta_no_provider_config_id_falls_back_to_default(self): + pmd = ProviderMetaData( + id="default", model=None, type="test_type", provider_type=ProviderType.EMBEDDING + ) + with patch( + "astrbot.core.provider.provider.provider_cls_map", + {"test_type": pmd}, + ): + ap = AbstractProvider(provider_config={"type": "test_type"}) + meta = ap.meta() + assert meta.id == "default" + + def test_test_does_not_raise(self): + ap = AbstractProvider(provider_config={"type": "test"}) + # test() is a no-op on AbstractProvider + ap.test() # should not raise + + def test_constructor_sets_provider_config(self): + config = {"type": "myprovider", "key": ["sk-abc"]} + ap = AbstractProvider(provider_config=config) + assert ap.provider_config is config + + +# ========================================================================= +# Provider (Chat) +# ========================================================================= + + +class _ConcreteProvider(Provider): + """Minimal concrete subclass for testing Provider abstract methods.""" + + def get_current_key(self) -> str: + return "key_override" + + def set_key(self, key: str) -> None: + self._key = key + + async def get_models(self) -> list[str]: + return ["model-a", "model-b"] + + async def text_chat(self, **kwargs) -> LLMResponse: + return LLMResponse(role="assistant", completion_text="mock reply") + + +class TestProvider: + def test_construction(self): + p = _ConcreteProvider( + provider_config={"type": "test", "key": ["sk-abc"]}, + provider_settings={}, + ) + assert p.provider_config["type"] == "test" + assert p.provider_settings == {} + + def test_get_current_key_abstract(self): + p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={}) + assert p.get_current_key() == "key_override" + + def test_get_keys_default(self): + p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={}) + assert p.get_keys() == [""] + + def test_get_keys_from_config(self): + p = _ConcreteProvider( + provider_config={"type": "test", "key": ["sk-1", "sk-2"]}, + provider_settings={}, + ) + assert p.get_keys() == ["sk-1", "sk-2"] + + def test_get_keys_none(self): + p = _ConcreteProvider( + provider_config={"type": "test", "key": None}, + provider_settings={}, + ) + assert p.get_keys() == [""] + + def test_set_key(self): + p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={}) + p.set_key("new-key") + assert p._key == "new-key" + + def test_get_models(self): + p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={}) + models = p.get_models() + assert models == ["model-a", "model-b"] + + def test_text_chat(self): + p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={}) + import asyncio + + resp = asyncio.run(p.text_chat(prompt="hi")) + assert isinstance(resp, LLMResponse) + assert resp.role == "assistant" + assert resp.completion_text == "mock reply" + + def test_text_chat_abstract_prevents_instantiation(self): + """Provider subclasses must override text_chat; can't instantiate without it.""" + class IncompleteProvider(Provider): + def get_current_key(self) -> str: + return "" + + def set_key(self, key: str) -> None: + pass + + async def get_models(self) -> list[str]: + return [] + + with pytest.raises(TypeError): + IncompleteProvider(provider_config={}, provider_settings={}) + + def test_text_chat_stream_raises_not_implemented(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + with pytest.raises(NotImplementedError): + import asyncio + + async def consume(): + gen = p.text_chat_stream(prompt="hi") + async for _ in gen: + pass + + asyncio.run(consume()) + + def test_pop_record_removes_non_system(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + ctx = [ + {"role": "system", "content": "You are a bot"}, + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "how are you"}, + ] + p.pop_record(ctx) + assert len(ctx) == 2 + assert ctx[0]["role"] == "system" + assert ctx[1]["role"] == "user" + assert ctx[1]["content"] == "how are you" + + def test_pop_record_with_no_system(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + ctx = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply"}, + ] + p.pop_record(ctx) + # Both should be popped + assert len(ctx) == 0 + + def test_pop_record_capped_at_two(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + ctx = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + {"role": "user", "content": "u3"}, + ] + p.pop_record(ctx) + assert len(ctx) == 4 + # system kept + last 3 non-system (u2, a2, u3) → but pop_record removes first 2 non-system + assert ctx[0] == {"role": "system", "content": "sys"} + assert ctx[1] == {"role": "user", "content": "u2"} + + @pytest.mark.asyncio + async def test_test(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + with patch.object(p, "text_chat", AsyncMock(return_value=LLMResponse(role="assistant"))): + await p.test(test_timeout=5.0) + + def test_ensure_message_to_dicts_none(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + result = p._ensure_message_to_dicts(None) + assert result == [] + + def test_ensure_message_to_dicts_empty(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + result = p._ensure_message_to_dicts([]) + assert result == [] + + def test_ensure_message_to_dicts_skips_checkpoint(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + checkpoint = {"role": "user", "content": "check", "checkpoint": True} + normal = {"role": "user", "content": "normal"} + with patch( + "astrbot.core.provider.provider.is_checkpoint_message", + side_effect=lambda m: m.get("checkpoint", False), + ): + result = p._ensure_message_to_dicts([checkpoint, normal]) + assert len(result) == 1 + assert result[0] == {"role": "user", "content": "normal"} + + def test_ensure_message_to_dicts_converts_message_objects(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + msg = MagicMock(spec=Message) + msg.model_dump.return_value = {"role": "user", "content": "from pydantic"} + result = p._ensure_message_to_dicts([msg]) + assert result == [{"role": "user", "content": "from pydantic"}] + + def test_ensure_message_to_dicts_passes_through_dicts(self): + p = _ConcreteProvider(provider_config={}, provider_settings={}) + d = {"role": "user", "content": "plain dict"} + result = p._ensure_message_to_dicts([d]) + assert result == [d] + + +# ========================================================================= +# STTProvider +# ========================================================================= + + +class _ConcreteSTTProvider(STTProvider): + async def get_text(self, audio_url: str) -> str: + return "transcribed text" + + +class TestSTTProvider: + def test_construction(self): + p = _ConcreteSTTProvider( + provider_config={"type": "stt_test"}, + provider_settings={}, + ) + assert p.provider_config["type"] == "stt_test" + + def test_get_text(self): + p = _ConcreteSTTProvider(provider_config={}, provider_settings={}) + import asyncio + + text = asyncio.run(p.get_text("/path/to/audio.wav")) + assert text == "transcribed text" + + def test_get_text_abstract(self): + class IncompleteSTT(STTProvider): + pass + + with pytest.raises(TypeError): + IncompleteSTT(provider_config={}, provider_settings={}) + + +# ========================================================================= +# TTSProvider +# ========================================================================= + + +class _ConcreteTTSProvider(TTSProvider): + async def get_audio(self, text: str) -> str: + return "/tmp/test_output.wav" + + +class TestTTSProvider: + def test_construction(self): + p = _ConcreteTTSProvider(provider_config={"type": "tts_test"}, provider_settings={}) + assert p.provider_config["type"] == "tts_test" + + def test_get_audio(self): + p = _ConcreteTTSProvider(provider_config={}, provider_settings={}) + import asyncio + + path = asyncio.run(p.get_audio("hello")) + assert path == "/tmp/test_output.wav" + + def test_support_stream_default(self): + p = _ConcreteTTSProvider(provider_config={}, provider_settings={}) + assert p.support_stream() is False + + def test_get_audio_stream_default_implementation(self): + p = _ConcreteTTSProvider(provider_config={}, provider_settings={}) + import asyncio + + text_queue: asyncio.Queue[str | None] = asyncio.Queue() + audio_queue: asyncio.Queue = asyncio.Queue() + + async def run_stream(): + # Send some text, then None to signal end + await text_queue.put("hello ") + await text_queue.put("world") + await text_queue.put(None) + await p.get_audio_stream(text_queue, audio_queue) + + asyncio.run(run_stream()) + result = asyncio.run(audio_queue.get()) + assert result is not None + text_part, audio_data = result + assert text_part == "hello world" + assert isinstance(audio_data, bytes) + # The None sentinel should follow + end = asyncio.run(audio_queue.get()) + assert end is None + + +# ========================================================================= +# EmbeddingProvider +# ========================================================================= + + +class _ConcreteEmbeddingProvider(EmbeddingProvider): + async def get_embedding(self, text: str) -> list[float]: + return [0.1, 0.2, 0.3] + + async def get_embeddings(self, text: list[str]) -> list[list[float]]: + return [[0.1, 0.2, 0.3] for _ in text] + + def get_dim(self) -> int: + return 3 + + +class TestEmbeddingProvider: + def test_construction(self): + p = _ConcreteEmbeddingProvider( + provider_config={"type": "emb_test"}, + provider_settings={}, + ) + assert p.provider_config["type"] == "emb_test" + + def test_get_embedding(self): + p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={}) + import asyncio + + emb = asyncio.run(p.get_embedding("hello")) + assert emb == [0.1, 0.2, 0.3] + + def test_get_embeddings(self): + p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={}) + import asyncio + + embs = asyncio.run(p.get_embeddings(["a", "b"])) + assert len(embs) == 2 + + def test_get_dim(self): + p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={}) + assert p.get_dim() == 3 + + def test_get_embeddings_batch_single_batch(self): + p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={}) + import asyncio + + embs = asyncio.run(p.get_embeddings_batch(["hello", "world"], batch_size=10)) + assert len(embs) == 2 + + def test_get_embeddings_batch_multiple_batches(self): + p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={}) + import asyncio + + texts = [f"text_{i}" for i in range(5)] + embs = asyncio.run(p.get_embeddings_batch(texts, batch_size=2, tasks_limit=5)) + assert len(embs) == 5 + + def test_get_embeddings_batch_with_progress_callback(self): + p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={}) + import asyncio + + progress = AsyncMock() + texts = [f"t{i}" for i in range(3)] + embs = asyncio.run( + p.get_embeddings_batch(texts, batch_size=2, tasks_limit=5, progress_callback=progress) + ) + assert len(embs) == 3 + assert progress.await_count >= 1 + + +# ========================================================================= +# RerankProvider +# ========================================================================= + + +class _ConcreteRerankProvider(RerankProvider): + async def rerank( + self, + query: str, + documents: list[str], + top_n: int | None = None, + ): + from astrbot.core.provider.entities import RerankResult + + return [ + RerankResult(index=0, relevance_score=0.95), + RerankResult(index=1, relevance_score=0.80), + ][:top_n] + + +class TestRerankProvider: + def test_construction(self): + p = _ConcreteRerankProvider( + provider_config={"type": "rerank_test"}, + provider_settings={}, + ) + assert p.provider_config["type"] == "rerank_test" + + def test_rerank(self): + p = _ConcreteRerankProvider(provider_config={}, provider_settings={}) + import asyncio + + results = asyncio.run(p.rerank("test query", ["doc1", "doc2"])) + assert len(results) == 2 + assert results[0].index == 0 + assert results[0].relevance_score == 0.95 + + def test_rerank_with_top_n(self): + p = _ConcreteRerankProvider(provider_config={}, provider_settings={}) + import asyncio + + results = asyncio.run(p.rerank("test query", ["doc1", "doc2", "doc3"], top_n=1)) + assert len(results) == 1 diff --git a/tests/unit/test_provider_register.py b/tests/unit/test_provider_register.py new file mode 100644 index 000000000..859684213 --- /dev/null +++ b/tests/unit/test_provider_register.py @@ -0,0 +1,244 @@ +"""Tests for astrbot.core.provider.register. + +Covers register_provider_adapter decorator, provider_registry list, +provider_cls_map dict, llm_tools import, duplicate registration, +and default config template processing. +""" + +import pytest + +from astrbot.core.provider.entities import ProviderMetaData, ProviderType +from astrbot.core.provider.register import ( + llm_tools, + provider_cls_map, + provider_registry, + register_provider_adapter, +) + + +# ========================================================================= +# Fixtures — fresh state per test +# ========================================================================= + + +@pytest.fixture(autouse=True) +def _clear_registries(): + """Clear global registries before and after each test to avoid cross-test pollution.""" + before_keys = set(provider_cls_map.keys()) + before_len = len(provider_registry) + + yield + + # Restore: remove any keys/items added during the test + added_keys = set(provider_cls_map.keys()) - before_keys + for k in added_keys: + provider_cls_map.pop(k, None) + + added_items = len(provider_registry) - before_len + for _ in range(added_items): + if provider_registry: + provider_registry.pop() + + +# ========================================================================= +# register_provider_adapter +# ========================================================================= + + +class TestRegisterProviderAdapter: + def test_basic_registration(self): + @register_provider_adapter("test_provider", "A test provider") + class FakeProvider: + pass + + assert "test_provider" in provider_cls_map + pmd = provider_cls_map["test_provider"] + assert isinstance(pmd, ProviderMetaData) + assert pmd.type == "test_provider" + assert pmd.desc == "A test provider" + assert pmd.cls_type is FakeProvider + assert pmd.provider_type == ProviderType.CHAT_COMPLETION # default + # Should also be in the registry list + assert pmd in provider_registry + + def test_registration_with_custom_provider_type(self): + @register_provider_adapter( + "stt_provider", + "STT provider", + provider_type=ProviderType.SPEECH_TO_TEXT, + ) + class FakeSTT: + pass + + pmd = provider_cls_map["stt_provider"] + assert pmd.provider_type == ProviderType.SPEECH_TO_TEXT + + def test_registration_with_display_name(self): + @register_provider_adapter( + "disp_provider", + "Display test", + provider_display_name="My Display Name", + ) + class DispProvider: + pass + + pmd = provider_cls_map["disp_provider"] + assert pmd.provider_display_name == "My Display Name" + + def test_registration_with_default_config_tmpl(self): + @register_provider_adapter( + "tmpl_provider", + "Template test", + default_config_tmpl={"key1": "val1"}, + ) + class TmplProvider: + pass + + pmd = provider_cls_map["tmpl_provider"] + assert pmd.default_config_tmpl is not None + # The decorator adds mandatory fields + assert pmd.default_config_tmpl["type"] == "tmpl_provider" + assert pmd.default_config_tmpl["enable"] is False + assert pmd.default_config_tmpl["id"] == "tmpl_provider" + assert pmd.default_config_tmpl["key1"] == "val1" + + def test_default_config_tmpl_preserves_existing_type(self): + @register_provider_adapter( + "tmpl2", + "desc", + default_config_tmpl={"type": "custom_type", "custom": True}, + ) + class T2: + pass + + pmd = provider_cls_map["tmpl2"] + # should NOT override existing type + assert pmd.default_config_tmpl["type"] == "custom_type" + assert pmd.default_config_tmpl["enable"] is False + assert pmd.default_config_tmpl["id"] == "tmpl2" + assert pmd.default_config_tmpl["custom"] is True + + def test_default_config_tmpl_preserves_existing_enable(self): + @register_provider_adapter( + "tmpl3", "desc", default_config_tmpl={"enable": True, "extra": "x"} + ) + class T3: + pass + + pmd = provider_cls_map["tmpl3"] + assert pmd.default_config_tmpl["enable"] is True + assert pmd.default_config_tmpl["type"] == "tmpl3" + assert pmd.default_config_tmpl["extra"] == "x" + + def test_default_config_tmpl_preserves_existing_id(self): + @register_provider_adapter( + "tmpl4", + "desc", + default_config_tmpl={"id": "custom_id"}, + ) + class T4: + pass + + pmd = provider_cls_map["tmpl4"] + assert pmd.default_config_tmpl["id"] == "custom_id" + + def test_default_config_tmpl_none(self): + @register_provider_adapter("no_tmpl", "No template") + class NoTmpl: + pass + + pmd = provider_cls_map["no_tmpl"] + assert pmd.default_config_tmpl is None + + def test_duplicate_registration_raises(self): + @register_provider_adapter("dup_provider", "First") + class FirstProvider: + pass + + with pytest.raises(ValueError, match="已经注册"): + @register_provider_adapter("dup_provider", "Second") + class SecondProvider: + pass + + # The original registration should remain intact + assert provider_cls_map["dup_provider"].desc == "First" + assert provider_cls_map["dup_provider"].cls_type is FirstProvider + + def test_multiple_registrations(self): + @register_provider_adapter("p1", "First provider") + class P1: + pass + + @register_provider_adapter("p2", "Second provider") + class P2: + pass + + assert len(provider_cls_map) == 2 + assert provider_cls_map["p1"].cls_type is P1 + assert provider_cls_map["p2"].cls_type is P2 + assert len(provider_registry) >= 2 + + def test_decorator_returns_the_class(self): + @register_provider_adapter("return_test", "Check return") + class ReturnClass: + pass + + # The decorator should return the class unchanged so it can be used normally + instance = ReturnClass() + assert isinstance(instance, ReturnClass) + + +# ========================================================================= +# llm_tools / FuncCall +# ========================================================================= + + +class TestLlmTools: + def test_llm_tools_is_importable(self): + from astrbot.core.provider.register import llm_tools as lt + + assert lt is not None + + def test_llm_tools_function_tool_manager_type(self): + from astrbot.core.provider.func_tool_manager import FunctionToolManager + + assert isinstance(llm_tools, FunctionToolManager) + + def test_llm_tools_starts_empty(self): + assert llm_tools.empty() is True + + def test_llm_tools_add_and_get(self): + async def fake_handler(**kwargs): + return "done" + + llm_tools.add_func( + name="test_func", + func_args=[{"name": "arg1", "type": "string", "description": "An arg"}], + desc="A test function", + handler=fake_handler, + ) + tool = llm_tools.get_func("test_func") + assert tool is not None + assert tool.name == "test_func" + llm_tools.remove_func("test_func") + assert llm_tools.get_func("test_func") is None + + def test_llm_tools_remove_nonexistent(self): + """Removing a function that does not exist should not raise.""" + llm_tools.remove_func("nonexistent_tool") # should not raise + + +# ========================================================================= +# Registry list / map invariants +# ========================================================================= + + +class TestRegistryInvariants: + def test_meta_id_is_default_in_registry(self): + @register_provider_adapter("invariant_test", "check id") + class InvProvider: + pass + + pmd = provider_cls_map["invariant_test"] + assert pmd.id == "default" + assert pmd.model is None diff --git a/tests/unit/test_recursive_chunker.py b/tests/unit/test_recursive_chunker.py new file mode 100644 index 000000000..3d2e64570 --- /dev/null +++ b/tests/unit/test_recursive_chunker.py @@ -0,0 +1,375 @@ +""" +Unit tests for RecursiveCharacterChunker. + +Covers construction, chunk method with various inputs (empty text, short text, +newline separators, character-level fallback, recursive splitting, custom +separators, and edge cases with chunk_size/overlap validation). +All tests isolate the chunker from any external dependencies. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from astrbot.core.knowledge_base.chunking.recursive import ( + RecursiveCharacterChunker, +) + + +# --------------------------------------------------------------- +# Construction +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerConstruction: + """Test construction of RecursiveCharacterChunker.""" + + def test_default_construction(self): + """Test default parameters are set correctly.""" + chunker = RecursiveCharacterChunker() + assert chunker.chunk_size == 500 + assert chunker.chunk_overlap == 100 + assert chunker.length_function is len + assert chunker.is_separator_regex is False + assert "\n\n" in chunker.separators + assert "" in chunker.separators # character fallback + + def test_custom_construction(self): + """Test custom parameters are applied.""" + chunker = RecursiveCharacterChunker( + chunk_size=1000, + chunk_overlap=200, + length_function=lambda x: len(x.split()), + is_separator_regex=True, + separators=["\n", " "], + ) + assert chunker.chunk_size == 1000 + assert chunker.chunk_overlap == 200 + assert chunker.length_function("hello world") == 2 + assert chunker.is_separator_regex is True + assert chunker.separators == ["\n", " "] + + def test_injectable_length_function(self): + """Test injecting a custom length function.""" + chunker = RecursiveCharacterChunker( + chunk_size=3, + length_function=lambda x: len(x.split()), + ) + assert chunker.chunk_size == 3 + assert chunker.length_function("a b c d") == 4 + + +# --------------------------------------------------------------- +# chunk() - basic cases +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerChunkBasic: + """Test basic chunk() behavior.""" + + @pytest.mark.asyncio + async def test_empty_text_returns_empty_list(self): + """Test that empty text returns an empty list.""" + chunker = RecursiveCharacterChunker() + result = await chunker.chunk("") + assert result == [] + + @pytest.mark.asyncio + async def test_whitespace_only_text(self): + """Test that whitespace-only text is handled.""" + chunker = RecursiveCharacterChunker() + result = await chunker.chunk(" \n\n ") + # Depends on separator logic; should not crash + assert isinstance(result, list) + + @pytest.mark.asyncio + async def test_short_text_returns_single_chunk(self): + """Test that text shorter than chunk_size returns single chunk.""" + chunker = RecursiveCharacterChunker(chunk_size=1000) + text = "Short text." + result = await chunker.chunk(text) + assert result == [text] + + @pytest.mark.asyncio + async def test_text_equal_to_chunk_size(self): + """Test that text exactly chunk_size returns single chunk.""" + chunker = RecursiveCharacterChunker(chunk_size=10) + text = "0123456789" + result = await chunker.chunk(text) + assert result == [text] + + +# --------------------------------------------------------------- +# chunk() - separator-driven splitting +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerSeparator: + """Test chunk() behavior with separator-based splitting.""" + + @pytest.mark.asyncio + async def test_splits_by_double_newline(self): + """Test that double newline is the preferred separator.""" + chunker = RecursiveCharacterChunker(chunk_size=10, chunk_overlap=0) + text = "para1\n\npara2\n\npara3" + result = await chunker.chunk(text) + # Each paragraph is <= chunk_size, so each should be a separate chunk + assert len(result) >= 1 + # All chunks should be non-empty strings + assert all(isinstance(c, str) and c for c in result) + + @pytest.mark.asyncio + async def test_splits_by_newline_when_double_newline_not_found(self): + """Test fallback to single newline when double newline is absent.""" + chunker = RecursiveCharacterChunker( + chunk_size=50, + chunk_overlap=0, + separators=["\n\n", "\n"], + ) + text = "line1\nline2\nline3" + result = await chunker.chunk(text) + # Without overlap and small lines, each line should be separate if lines fit + assert isinstance(result, list) + assert all(isinstance(c, str) for c in result) + + @pytest.mark.asyncio + async def test_falls_back_to_character_splitting(self): + """Test that chunker falls back to character splitting when no separator matches.""" + chunker = RecursiveCharacterChunker( + chunk_size=5, + chunk_overlap=0, + separators=["\n\n", ""], + ) + text = "abcdefghij" + result = await chunker.chunk(text) + # Character splitting: step = chunk_size - overlap = 5 + assert result == ["abcde", "fghij"] + + @pytest.mark.asyncio + async def test_custom_separator(self): + """Test that a custom separator is used for splitting.""" + chunker = RecursiveCharacterChunker( + chunk_size=100, + chunk_overlap=0, + separators=["|"], + ) + text = "part1|part2|part3" + result = await chunker.chunk(text) + # Each part is well under chunk_size, so each should be a separate chunk + assert len(result) == 3 + assert result[0] == "part1|" + assert result[1] == "part2|" + assert result[2] == "part3" + + +# --------------------------------------------------------------- +# chunk() - overlap behavior +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerOverlap: + """Test overlap behavior in chunk().""" + + @pytest.mark.asyncio + async def test_character_split_with_overlap(self): + """Test that character-level split respects overlap.""" + chunker = RecursiveCharacterChunker( + chunk_size=6, + chunk_overlap=2, + separators=[""], + ) + text = "abcdefghij" + result = await chunker.chunk(text) + # step = 6 - 2 = 4, so: "abcdef", "efghij" + assert result == ["abcdef", "efghij"] + + @pytest.mark.asyncio + async def test_overlap_from_kwargs_overrides_instance_default(self): + """Test that kwargs chunk_overlap overrides the instance default.""" + chunker = RecursiveCharacterChunker( + chunk_size=6, + chunk_overlap=0, + separators=[""], + ) + text = "abcdefghij" + result = await chunker.chunk(text, chunk_overlap=2) + # step = 6 - 2 = 4 + assert result == ["abcdef", "efghij"] + + @pytest.mark.asyncio + async def test_chunk_size_from_kwargs_overrides_instance_default(self): + """Test that kwargs chunk_size overrides the instance default.""" + chunker = RecursiveCharacterChunker( + chunk_size=500, + chunk_overlap=0, + separators=[""], + ) + text = "abcdefghij" + result = await chunker.chunk(text, chunk_size=5) + # step = 5 - 0 = 5 + assert result == ["abcde", "fghij"] + + +# --------------------------------------------------------------- +# chunk() - recursive splitting of oversized segments +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerRecursive: + """Test recursive splitting of segments that exceed chunk_size.""" + + @pytest.mark.asyncio + async def test_recursive_split_oversized_segment(self): + """Test that a single segment larger than chunk_size is recursively split.""" + chunker = RecursiveCharacterChunker( + chunk_size=10, + chunk_overlap=0, + separators=["\n\n", ""], + ) + # Single paragraph (no double newline) that exceeds chunk_size + text = "a" * 25 + result = await chunker.chunk(text) + # Should split into 3 character-level chunks: 10 + 10 + 5 + assert len(result) == 3 + assert result[0] == "a" * 10 + assert result[1] == "a" * 10 + assert result[2] == "a" * 5 + + @pytest.mark.asyncio + async def test_recursive_and_normal_chunks_mixed(self): + """Test mixing normal chunks and recursively split oversized segments.""" + chunker = RecursiveCharacterChunker( + chunk_size=10, + chunk_overlap=0, + separators=["\n\n", ""], + ) + # First paragraph fits, second paragraph is oversized, third fits + text = "short" + "\n\n" + ("a" * 25) + "\n\n" + "tiny" + result = await chunker.chunk(text) + # short\n\n + a*10 + a*10 + a*5 + \n\n + tiny -> should be 4 chunks + # Note: the separator is included in each split + assert len(result) >= 3 + assert all(isinstance(c, str) and c for c in result) + + +# --------------------------------------------------------------- +# _split_by_character +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerSplitByCharacter: + """Test the _split_by_character private method.""" + + def test_split_by_character_defaults(self): + """Test _split_by_character uses instance defaults.""" + chunker = RecursiveCharacterChunker(chunk_size=4, chunk_overlap=1) + result = chunker._split_by_character("abcdefgh") + assert result == ["abcd", "defg", "gh"] + + def test_split_by_character_explicit_params(self): + """Test _split_by_character with explicit chunk_size and overlap.""" + chunker = RecursiveCharacterChunker() + result = chunker._split_by_character("abcdefgh", chunk_size=4, overlap=1) + assert result == ["abcd", "defg", "gh"] + + def test_split_by_character_long_exact_fit(self): + """Test split when text length divides evenly into chunk_size.""" + chunker = RecursiveCharacterChunker(chunk_size=4, chunk_overlap=0) + result = chunker._split_by_character("abcdefgh") + assert result == ["abcd", "efgh"] + + def test_split_by_character_short_text(self): + """Test split when text is shorter than chunk_size.""" + chunker = RecursiveCharacterChunker(chunk_size=100) + result = chunker._split_by_character("hello") + assert result == ["hello"] + + def test_split_by_character_raises_on_zero_chunk_size(self): + """Test that chunk_size <= 0 raises ValueError.""" + chunker = RecursiveCharacterChunker() + with pytest.raises(ValueError, match="chunk_size must be greater than 0"): + chunker._split_by_character("test", chunk_size=0) + + def test_split_by_character_raises_on_negative_overlap(self): + """Test that negative overlap raises ValueError.""" + chunker = RecursiveCharacterChunker() + with pytest.raises(ValueError, match="chunk_overlap must be non-negative"): + chunker._split_by_character("test", chunk_overlap=-1) + + def test_split_by_character_raises_on_overlap_ge_chunk_size(self): + """Test that overlap >= chunk_size raises ValueError.""" + chunker = RecursiveCharacterChunker() + with pytest.raises(ValueError, match="chunk_overlap must be less than chunk_size"): + chunker._split_by_character("test", chunk_size=5, chunk_overlap=5) + + def test_single_character_chunks(self): + """Test splitting with chunk_size=1.""" + chunker = RecursiveCharacterChunker(chunk_size=1, chunk_overlap=0) + result = chunker._split_by_character("abc") + assert result == ["a", "b", "c"] + + def test_overlap_equal_chunk_size_minus_one(self): + """Test with maximum valid overlap.""" + chunker = RecursiveCharacterChunker(chunk_size=5, chunk_overlap=4) + result = chunker._split_by_character("abcdefghij") + # step = 1, each chunk slides by 1 char + assert result == ["abcde", "bcdef", "cdefg", "defgh", "efghi", "fghij"] + + +# --------------------------------------------------------------- +# chunk() - integration scenarios +# --------------------------------------------------------------- + + +class TestRecursiveCharacterChunkerIntegration: + """Integration-style tests combining multiple logic paths.""" + + @pytest.mark.asyncio + async def test_paragraphs_with_oversized_lines(self): + """Test a realistic paragraph mix with some oversized lines.""" + chunker = RecursiveCharacterChunker(chunk_size=20, chunk_overlap=0) + lines = ["A" * 30, "short", "B" * 30] + text = "\n".join(lines) + result = await chunker.chunk(text) + # Should split the oversized lines but keep short lines as-is + assert all(isinstance(c, str) and c for c in result) + assert len(result) >= 3 + + @pytest.mark.asyncio + async def test_separator_in_text_but_not_at_separator_boundary(self): + """Test that chunker finds the best separator when multiple are present.""" + chunker = RecursiveCharacterChunker( + chunk_size=10, + chunk_overlap=0, + separators=["\n\n", "\n", ""], + ) + # Has both double-newline and single-newline; prefer double-newline + text = "aaaa\n\nbbbb\ncccc" + result = await chunker.chunk(text) + # Should split on \n\n first, then possibly \n for remaining + assert isinstance(result, list) + + @pytest.mark.asyncio + async def test_text_preserves_ordering(self): + """Test that chunks preserve the original text ordering.""" + chunker = RecursiveCharacterChunker(chunk_size=10, chunk_overlap=0) + text = "first\n\nsecond\n\nthird\n\nfourth" + result = await chunker.chunk(text) + # Concatenation of all chunks should contain all words in order + combined = "".join(result) + for word in ["first", "second", "third", "fourth"]: + assert word in combined + + @pytest.mark.asyncio + async def test_no_duplicate_chunks(self): + """Test that chunker does not produce duplicate content.""" + chunker = RecursiveCharacterChunker( + chunk_size=10, + chunk_overlap=0, + separators=[""], + ) + text = "abcdefghij" + result = await chunker.chunk(text) + combined = "".join(result) + # Combined length should equal original length (no overlap) + assert len(combined) == len(text) diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py new file mode 100644 index 000000000..c79e1a25e --- /dev/null +++ b/tests/unit/test_registry.py @@ -0,0 +1,331 @@ +"""Unit tests for astrbot.core.tools.registry: builtin_tool decorator, tool registration.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from pydantic.dataclasses import dataclass + +from astrbot.core.agent.tool import FunctionTool +from astrbot.core.tools.registry import ( + BuiltinToolConfigCondition, + BuiltinToolConfigRule, + _BUILTIN_TOOL_CONFIG_RULES, + _builtin_tool_classes_by_name, + _builtin_tool_names_by_class, + _get_config_value, + _json_safe, + _MISSING, + _resolve_builtin_tool_name, + builtin_tool, + ensure_builtin_tools_loaded, + get_builtin_tool_class, + get_builtin_tool_config_rule, + get_builtin_tool_config_statuses, + get_builtin_tool_config_tags, + get_builtin_tool_name, + iter_builtin_tool_classes, +) + + +@pytest.fixture(autouse=True) +def _clean_registry(): + """Clean the global builtin tool registry before and after each test.""" + before_classes = dict(_builtin_tool_classes_by_name) + before_names = dict(_builtin_tool_names_by_class) + before_rules = dict(_BUILTIN_TOOL_CONFIG_RULES) + yield + _builtin_tool_classes_by_name.clear() + _builtin_tool_classes_by_name.update(before_classes) + _builtin_tool_names_by_class.clear() + _builtin_tool_names_by_class.update(before_names) + _BUILTIN_TOOL_CONFIG_RULES.clear() + _BUILTIN_TOOL_CONFIG_RULES.update(before_rules) + + +class TestBuiltinToolDecorator: + """builtin_tool decorator registration.""" + + def test_decorator_without_call_registers(self): + """Using @builtin_tool without parentheses registers the class.""" + + @builtin_tool + class MyTool(FunctionTool): + name: str = "my_tool" + description: str = "My custom tool" + + assert get_builtin_tool_class("my_tool") is MyTool + assert get_builtin_tool_name(MyTool) == "my_tool" + + def test_decorator_with_call_registers(self): + """Using @builtin_tool() with parentheses registers the class.""" + + @builtin_tool() + class AnotherTool(FunctionTool): + name: str = "another_tool" + description: str = "Another tool" + + assert get_builtin_tool_class("another_tool") is AnotherTool + + def test_decorator_with_config(self): + """@builtin_tool(config={...}) registers with config rules.""" + + @builtin_tool(config={"feature.enabled": True, "mode": ("auto", "manual")}) + class ConfigTool(FunctionTool): + name: str = "config_tool" + description: str = "Tool with config" + + rule = get_builtin_tool_config_rule("config_tool") + assert rule is not None + assert len(rule.conditions) == 2 + + def test_name_conflict_raises(self): + """Registering the same name with a different class raises ValueError.""" + + @builtin_tool + class First(FunctionTool): + name: str = "conflict_tool" + description: str = "first" + + with pytest.raises(ValueError, match="name conflict"): + + @builtin_tool + class Second(FunctionTool): + name: str = "conflict_tool" + description: str = "second" + + def test_same_class_no_conflict(self): + """Registering the same class again with the same name does not raise.""" + + @builtin_tool + class SameTool(FunctionTool): + name: str = "same_tool" + description: str = "same" + + # Registering the same class again should not raise + builtin_tool(SameTool) + assert get_builtin_tool_class("same_tool") is SameTool + + def test_resolve_tool_name_from_field(self): + """_resolve_builtin_tool_name reads from the 'name' dataclass field.""" + + @dataclass + class ResolveMe: + name: str = "resolved_name" + + # Since ResolveMe is not a FunctionTool, we need to test _resolve_builtin_tool_name + # by temporarily making it look like a FunctionTool subclass + name = _resolve_builtin_tool_name(ResolveMe) + assert name == "resolved_name" + + def test_resolve_tool_name_raises_when_missing(self): + """_resolve_builtin_tool_name raises ValueError when no name is found.""" + + class NoName: + pass + + with pytest.raises(ValueError, match="does not define a valid name"): + _resolve_builtin_tool_name(NoName) + + +class TestGetAndIter: + """Query functions for the builtin tool registry.""" + + def test_get_nonexistent_class_returns_none(self): + """get_builtin_tool_class returns None for unknown names.""" + assert get_builtin_tool_class("nonexistent_tool") is None + + def test_get_nonexistent_name_returns_none(self): + """get_builtin_tool_name returns None for unknown classes.""" + class Random(FunctionTool): + name: str = "random" + description: str = "r" + assert get_builtin_tool_name(Random) is None + + def test_iter_builtin_tool_classes_empty(self): + """iter_builtin_tool_classes returns empty tuple when nothing registered.""" + classes = iter_builtin_tool_classes() + # Only pre-existing builtins may be present, but the fixture resets to original state. + # We just check it's a tuple. + assert isinstance(classes, tuple) + + def test_iter_after_registration(self): + """iter_builtin_tool_classes includes newly registered tools.""" + + @builtin_tool + class IterTool(FunctionTool): + name: str = "iter_tool" + description: str = "iter" + + classes = iter_builtin_tool_classes() + assert IterTool in classes + + +class TestConfigCondition: + """BuiltinToolConfigCondition evaluation.""" + + def test_equals_condition_match(self): + """'equals' operator returns matched=True when values match.""" + cond = BuiltinToolConfigCondition(key="enabled", operator="equals", expected=True) + result = cond.evaluate({"enabled": True}) + assert result["matched"] is True + + def test_equals_condition_mismatch(self): + """'equals' operator returns matched=False when values differ.""" + cond = BuiltinToolConfigCondition(key="enabled", operator="equals", expected=True) + result = cond.evaluate({"enabled": False}) + assert result["matched"] is False + + def test_in_condition_match(self): + """'in' operator returns matched=True when value is in expected.""" + cond = BuiltinToolConfigCondition(key="mode", operator="in", expected=("a", "b", "c")) + result = cond.evaluate({"mode": "b"}) + assert result["matched"] is True + + def test_in_condition_mismatch(self): + """'in' operator returns matched=False when value is not in expected.""" + cond = BuiltinToolConfigCondition(key="mode", operator="in", expected=("a", "b")) + result = cond.evaluate({"mode": "c"}) + assert result["matched"] is False + + def test_truthy_condition_match(self): + """'truthy' operator returns matched=True for truthy values.""" + cond = BuiltinToolConfigCondition(key="timeout", operator="truthy") + result = cond.evaluate({"timeout": 30}) + assert result["matched"] is True + + def test_truthy_condition_mismatch(self): + """'truthy' operator returns matched=False for falsy values.""" + cond = BuiltinToolConfigCondition(key="timeout", operator="truthy") + result = cond.evaluate({"timeout": 0}) + assert result["matched"] is False + + def test_custom_condition(self): + """'custom' operator delegates to the expected field.""" + cond = BuiltinToolConfigCondition(key="custom_key", operator="custom", expected=True) + result = cond.evaluate({}) + assert result["matched"] is True + + def test_unsupported_operator_raises(self): + """An unknown operator raises ValueError.""" + cond = BuiltinToolConfigCondition(key="k", operator="bad_op") + with pytest.raises(ValueError, match="Unsupported builtin tool config operator"): + cond.evaluate({}) + + def test_missing_key_returns_missing(self): + """A key that is not present in config returns _MISSING as actual.""" + cond = BuiltinToolConfigCondition(key="missing.key", operator="truthy") + result = cond.evaluate({}) + assert result["actual"] is None + + def test_nested_key_access(self): + """_get_config_value traverses dot-separated keys.""" + config = {"a": {"b": {"c": 42}}} + assert _get_config_value(config, "a.b.c") == 42 + assert _get_config_value(config, "a.b.missing") is _MISSING + assert _get_config_value(config, "x") is _MISSING + + def test_json_safe_converts_tuple_to_list(self): + """_json_safe converts tuples to lists.""" + result = _json_safe((1, (2, 3))) + assert result == [1, [2, 3]] + + def test_json_safe_dict(self): + """_json_safe processes dicts recursively.""" + result = _json_safe({"a": (1, 2), "b": "hello"}) + assert result == {"a": [1, 2], "b": "hello"} + + +class TestBuiltinToolConfigRule: + """BuiltinToolConfigRule evaluation.""" + + def test_rule_with_conditions(self): + """A rule with conditions evaluates all of them.""" + c1 = BuiltinToolConfigCondition(key="x", operator="equals", expected=1) + c2 = BuiltinToolConfigCondition(key="y", operator="truthy") + rule = BuiltinToolConfigRule(conditions=(c1, c2)) + results = rule.evaluate({"x": 1, "y": True}) + assert len(results) == 2 + assert all(r["matched"] for r in results) + + def test_rule_with_evaluator(self): + """A rule with an evaluator callable uses it instead of conditions.""" + def my_evaluator(config): + return [{"key": "custom", "matched": True}] + rule = BuiltinToolConfigRule(evaluator=my_evaluator) + results = rule.evaluate({"anything": 1}) + assert results == [{"key": "custom", "matched": True}] + + def test_rule_conditions_are_frozen(self): + """BuiltinToolConfigRule and its conditions are frozen dataclasses.""" + rule = BuiltinToolConfigRule(conditions=()) + with pytest.raises(Exception): + rule.conditions = ("cannot", "change") + + +class TestGetBuiltinToolConfigStatuses: + """get_builtin_tool_config_statuses integration.""" + + def test_no_rule_returns_empty(self): + """Getting statuses for a tool with no config rule returns [].""" + statuses = get_builtin_tool_config_statuses("nonexistent", [{"config": {}}]) + assert statuses == [] + + def test_statuses_with_matching_config(self): + """Statuses are returned with enabled=True when all conditions match.""" + + @builtin_tool(config={"enabled": True}) + class StatusTool(FunctionTool): + name: str = "status_tool" + description: str = "test" + + entries = [{"conf_id": "1", "conf_name": "cfg1", "config": {"enabled": True}}] + statuses = get_builtin_tool_config_statuses("status_tool", entries) + assert len(statuses) == 1 + assert statuses[0]["enabled"] is True + + def test_statuses_with_non_matching_config(self): + """Statuses are returned with enabled=False when some conditions fail.""" + + @builtin_tool(config={"enabled": True}) + class StatusTool2(FunctionTool): + name: str = "status_tool2" + description: str = "test" + + entries = [{"conf_id": "2", "conf_name": "cfg2", "config": {"enabled": False}}] + statuses = get_builtin_tool_config_statuses("status_tool2", entries) + assert len(statuses) == 1 + assert statuses[0]["enabled"] is False + assert len(statuses[0]["failed_conditions"]) > 0 + + def test_get_tags_filters_enabled(self): + """get_builtin_tool_config_tags only returns entries where enabled is True.""" + + @builtin_tool(config={"enabled": True}) + class TagTool(FunctionTool): + name: str = "tag_tool" + description: str = "test" + + entries = [ + {"conf_id": "1", "conf_name": "on", "config": {"enabled": True}}, + {"conf_id": "2", "conf_name": "off", "config": {"enabled": False}}, + ] + tags = get_builtin_tool_config_tags("tag_tool", entries) + assert len(tags) == 1 + assert tags[0]["conf_id"] == "1" + + +class TestEnsureBuiltinToolsLoaded: + """ensure_builtin_tools_loaded idempotency.""" + + def test_load_is_idempotent(self): + """Calling ensure_builtin_tools_loaded twice does not raise.""" + # The first call may fail if the builtin modules have missing deps; catch. + try: + ensure_builtin_tools_loaded() + except Exception: + pass + # Second call should also not raise (just returns early). + ensure_builtin_tools_loaded() diff --git a/tests/unit/test_run_context.py b/tests/unit/test_run_context.py new file mode 100644 index 000000000..82f3c1262 --- /dev/null +++ b/tests/unit/test_run_context.py @@ -0,0 +1,91 @@ +"""Unit tests for astrbot.core.agent.run_context: ContextWrapper.""" + +from __future__ import annotations + +import pytest + +from astrbot.core.agent.run_context import ContextWrapper, NoContext +from astrbot.core.agent.message import Message + + +class TestContextWrapper: + """ContextWrapper construction and default values.""" + + def test_default_messages_is_empty_list(self): + """ContextWrapper starts with an empty messages list.""" + ctx = ContextWrapper(context="test") + assert ctx.messages == [] + + def test_default_tool_call_timeout(self): + """Default tool_call_timeout is 120 seconds.""" + ctx = ContextWrapper(context="test") + assert ctx.tool_call_timeout == 120 + + def test_context_string(self): + """context holds a string value.""" + ctx = ContextWrapper(context="hello_world") + assert ctx.context == "hello_world" + + def test_context_integer(self): + """context holds an integer value.""" + ctx = ContextWrapper(context=42) + assert ctx.context == 42 + + def test_context_dict(self): + """context holds a dict value.""" + data = {"key": "value", "num": 1} + ctx = ContextWrapper(context=data) + assert ctx.context == data + + def test_context_none(self): + """context holds None.""" + ctx = ContextWrapper(context=None) + assert ctx.context is None + + def test_messages_can_be_appended(self): + """messages list supports appending Message objects.""" + ctx = ContextWrapper(context="test") + msg = Message(role="user", content="hi") + ctx.messages.append(msg) + assert len(ctx.messages) == 1 + assert ctx.messages[0].content == "hi" + + def test_messages_can_be_replaced(self): + """messages field can be replaced with a new list.""" + msgs = [Message(role="user", content="a"), Message(role="assistant", content="b")] + ctx = ContextWrapper(context="test", messages=msgs) + assert len(ctx.messages) == 2 + + def test_custom_tool_call_timeout(self): + """tool_call_timeout can be customized.""" + ctx = ContextWrapper(context="test", tool_call_timeout=30) + assert ctx.tool_call_timeout == 30 + + def test_generic_type_parameter(self): + """ContextWrapper is generic; it accepts a type parameter.""" + ctx: ContextWrapper[str] = ContextWrapper(context="typed") + assert ctx.context == "typed" + + +class TestNoContext: + """NoContext is a ContextWrapper[None] singleton-like pattern.""" + + def test_no_context_is_contextwrapper(self): + """NoContext is an instance of ContextWrapper.""" + nc = NoContext() + assert isinstance(nc, ContextWrapper) + + def test_no_context_context_is_none(self): + """NoContext always carries context=None.""" + nc = NoContext() + assert nc.context is None + + def test_no_context_messages_empty(self): + """NoContext starts with empty messages.""" + nc = NoContext() + assert nc.messages == [] + + def test_no_context_default_timeout(self): + """NoContext has the default 120s timeout.""" + nc = NoContext() + assert nc.tool_call_timeout == 120 diff --git a/tests/unit/test_skill_manager.py b/tests/unit/test_skill_manager.py new file mode 100644 index 000000000..cb5724bd8 --- /dev/null +++ b/tests/unit/test_skill_manager.py @@ -0,0 +1,829 @@ +""" +Unit tests for SkillManager. + +Covers construction, list_skills with various combinations (active_only, +runtime modes, sandbox caching), set_sandbox_skills_cache, +get_sandbox_skills_cache_status, is_sandbox_only_skill, set_skill_active, +delete_skill, and config persistence. +All tests use mocks to isolate SkillManager from the filesystem and I/O. +""" + +import json +import os +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, mock_open, patch + +import pytest + +from astrbot.core.skills.skill_manager import ( + DEFAULT_SKILLS_CONFIG, + SANDBOX_SKILLS_CACHE_FILENAME, + SKILLS_CONFIG_FILENAME, + SANDBOX_SKILLS_ROOT, + SANDBOX_WORKSPACE_ROOT, + SkillInfo, + SkillManager, + _normalize_skill_name, + _normalize_cached_sandbox_skill_path, + _is_ignored_zip_entry, + _sanitize_prompt_path_for_prompt, + _sanitize_prompt_description, + _sanitize_skill_display_name, + build_skills_prompt, + _parse_frontmatter, +) + + +# --------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------- + + +@pytest.fixture +def mock_astrbot_paths(): + """Create mock AstrbotPaths with in-memory-like attributes.""" + paths = MagicMock() + paths.skills = "/tmp/.astrbot/skills" + paths.config = Path("/tmp/.astrbot/config") + paths.data = Path("/tmp/.astrbot/data") + paths.temp = Path("/tmp/.astrbot/temp") + return paths + + +@pytest.fixture +def skill_manager(mock_astrbot_paths): + """Create a SkillManager with mocked paths and no real FS side effects.""" + with ( + patch("os.makedirs") as mock_makedirs, + patch.object(Path, "iterdir", return_value=[]), + ): + mgr = SkillManager( + skills_root="/tmp/test_skills", + astrbot_paths=mock_astrbot_paths, + ) + return mgr + + +@pytest.fixture +def sandbox_skill_entry(): + """Create a canned sandbox skill entry as found in cache.""" + return { + "name": "sandbox-skill", + "description": "A sandbox preset skill", + "path": f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/sandbox-skill/SKILL.md", + } + + +# --------------------------------------------------------------- +# Construction +# --------------------------------------------------------------- + + +class TestSkillManagerConstruction: + """Test SkillManager construction.""" + + def test_construction_sets_attributes(self, mock_astrbot_paths): + """Test that construction sets expected attributes and creates directories.""" + with patch("os.makedirs") as mock_makedirs: + mgr = SkillManager( + skills_root="/tmp/test_skills", + astrbot_paths=mock_astrbot_paths, + ) + assert mgr.skills_root == "/tmp/test_skills" + assert mgr.astrbot_paths is mock_astrbot_paths + mock_makedirs.assert_called_once_with("/tmp/test_skills", exist_ok=True) + + def test_construction_default_skills_root(self, mock_astrbot_paths): + """Test that skills_root defaults to astrbot_paths.skills.""" + with patch("os.makedirs"): + mgr = SkillManager(astrbot_paths=mock_astrbot_paths) + assert mgr.skills_root == str(mock_astrbot_paths.skills) + + def test_construction_config_path(self, mock_astrbot_paths): + """Test that config_path is derived from astrbot_paths.""" + with patch("os.makedirs"): + mgr = SkillManager( + skills_root="/tmp/foo", + astrbot_paths=mock_astrbot_paths, + ) + expected = str(mock_astrbot_paths.config / SKILLS_CONFIG_FILENAME) + assert mgr.config_path == expected + + def test_construction_sandbox_cache_path(self, mock_astrbot_paths): + """Test that sandbox_skills_cache_path is derived correctly.""" + with patch("os.makedirs"): + mgr = SkillManager( + skills_root="/tmp/foo", + astrbot_paths=mock_astrbot_paths, + ) + expected = str( + mock_astrbot_paths.data / SANDBOX_SKILLS_CACHE_FILENAME, + ) + assert mgr.sandbox_skills_cache_path == expected + + +# --------------------------------------------------------------- +# list_skills - config loading +# --------------------------------------------------------------- + + +class TestSkillManagerListSkillsConfig: + """Test list_skills config loading behavior.""" + + def test_list_skills_creates_default_config_when_missing( + self, + skill_manager, + mock_astrbot_paths, + ): + """Test that list_skills creates a default config when config file is missing.""" + # Simulate config file does not exist + with ( + patch("os.path.exists", return_value=False) as mock_exists, + patch("builtins.open", mock_open()) as mock_file, + patch.object(Path, "iterdir", return_value=[]), + patch.object(Path, "exists", return_value=True), + ): + # exists is called multiple times; handle skills_root check too + mock_exists.side_effect = lambda p: False # all False + + with patch.object(Path, "mkdir"): + result = skill_manager.list_skills() + + assert result == [] + # Default config should have been saved + mock_file.assert_any_call( + skill_manager.config_path, + "w", + encoding="utf-8", + ) + + def test_list_skills_loads_existing_skills_from_config( + self, + skill_manager, + mock_astrbot_paths, + ): + """Test that list_skills reads skills from skill config files on disk.""" + config_data = json.dumps({"skills": {"my-skill": {"active": True}}}) + + skill_dir = MagicMock(spec=Path) + skill_dir.is_dir.return_value = True + skill_dir.name = "my-skill" + + skill_md = MagicMock(spec=Path) + skill_md.exists.return_value = True + skill_md.read_text.return_value = ( + "---\nname: my-skill\ndescription: My test skill\n---\nDo stuff." + ) + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=config_data)), + patch.object(Path, "iterdir", return_value=[skill_dir]), + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=skill_md, + ), + ): + result = skill_manager.list_skills() + + assert len(result) == 1 + assert result[0].name == "my-skill" + assert result[0].description == "My test skill" + assert result[0].active is True + assert result[0].local_exists is True + assert result[0].source_type == "local_only" + + def test_list_skills_skips_directories_without_skill_md( + self, + skill_manager, + ): + """Test that directories without SKILL.md are skipped.""" + config_data = json.dumps({"skills": {}}) + skill_dir = MagicMock(spec=Path) + skill_dir.is_dir.return_value = True + skill_dir.name = "not-a-skill" + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=config_data)), + patch.object(Path, "iterdir", return_value=[skill_dir]), + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=None, + ), + ): + result = skill_manager.list_skills() + + assert result == [] + + def test_list_skills_skips_non_directory_entries( + self, + skill_manager, + ): + """Test that non-directory entries in skills_root are skipped.""" + config_data = json.dumps({"skills": {}}) + file_entry = MagicMock(spec=Path) + file_entry.is_dir.return_value = False + file_entry.name = "README.md" + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=config_data)), + patch.object(Path, "iterdir", return_value=[file_entry]), + ): + result = skill_manager.list_skills() + + assert result == [] + + +# --------------------------------------------------------------- +# list_skills - active_only filter +# --------------------------------------------------------------- + + +class TestSkillManagerActiveOnly: + """Test the active_only filter in list_skills.""" + + def test_list_skills_active_only_excludes_inactive( + self, + skill_manager, + ): + """Test that active_only=True excludes inactive skills.""" + config_data = json.dumps( + { + "skills": { + "active-skill": {"active": True}, + "inactive-skill": {"active": False}, + }, + }, + ) + + active_md = MagicMock(spec=Path) + active_md.read_text.return_value = "---\nname: active-skill\ndescription: Active\n---" + inactive_md = MagicMock(spec=Path) + inactive_md.read_text.return_value = ( + "---\nname: inactive-skill\ndescription: Inactive\n---" + ) + + def normalize_skill_markdown_path(skill_dir): + if skill_dir.name == "active-skill": + return active_md + if skill_dir.name == "inactive-skill": + return inactive_md + return None + + active_dir = MagicMock(spec=Path) + active_dir.is_dir.return_value = True + active_dir.name = "active-skill" + + inactive_dir = MagicMock(spec=Path) + inactive_dir.is_dir.return_value = True + inactive_dir.name = "inactive-skill" + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=config_data)), + patch.object( + Path, + "iterdir", + return_value=[active_dir, inactive_dir], + ), + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + side_effect=normalize_skill_markdown_path, + ), + ): + result = skill_manager.list_skills(active_only=True) + + assert len(result) == 1 + assert result[0].name == "active-skill" + + def test_list_skills_active_only_false_returns_all( + self, + skill_manager, + ): + """Test that active_only=False returns all skills.""" + config_data = json.dumps( + { + "skills": { + "active-skill": {"active": True}, + "inactive-skill": {"active": False}, + }, + }, + ) + + def make_skill_dir(name, md): + d = MagicMock(spec=Path) + d.is_dir.return_value = True + d.name = name + return d + + def normalize_skill_markdown_path(skill_dir): + md = MagicMock(spec=Path) + md.read_text.return_value = f"---\nname: {skill_dir.name}\ndescription: desc\n---" + return md + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=config_data)), + patch.object( + Path, + "iterdir", + return_value=[ + make_skill_dir("active-skill", None), + make_skill_dir("inactive-skill", None), + ], + ), + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + side_effect=normalize_skill_markdown_path, + ), + ): + result = skill_manager.list_skills(active_only=False) + + assert len(result) == 2 + + +# --------------------------------------------------------------- +# list_skills - sandbox runtime +# --------------------------------------------------------------- + + +class TestSkillManagerSandbox: + """Test list_skills with runtime='sandbox'.""" + + def test_list_skills_sandbox_includes_sandbox_only_skills( + self, + skill_manager, + sandbox_skill_entry, + ): + """Test that sandbox runtime includes sandbox-only skills from cache.""" + config_data = json.dumps({"skills": {}}) + + cache_data = json.dumps( + { + "version": 1, + "skills": [sandbox_skill_entry], + }, + ) + + def exists_side_effect(path): + """Config file exists, sandbox cache exists.""" + return True + + def open_side_effect(path, *args, **kwargs): + if SANDBOX_SKILLS_CACHE_FILENAME in path: + return mock_open(read_data=cache_data).return_value + return mock_open(read_data=config_data).return_value + + with ( + patch("os.path.exists", side_effect=exists_side_effect), + patch("builtins.open", side_effect=open_side_effect), + patch.object(Path, "iterdir", return_value=[]), + ): + result = skill_manager.list_skills(runtime="sandbox") + + assert len(result) == 1 + assert result[0].name == "sandbox-skill" + assert result[0].source_type == "sandbox_only" + assert result[0].local_exists is False + assert result[0].sandbox_exists is True + + def test_list_skills_sandbox_marks_both_synced( + self, + skill_manager, + sandbox_skill_entry, + ): + """Test that skills existing both locally and in sandbox cache are 'both'.""" + config_data = json.dumps( + {"skills": {"local-skill": {"active": True}}}, + ) + + cache_data = json.dumps( + { + "version": 1, + "skills": [ + { + "name": "local-skill", + "description": "Synced skill", + "path": f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/local-skill/SKILL.md", + }, + ], + }, + ) + + skill_dir = MagicMock(spec=Path) + skill_dir.is_dir.return_value = True + skill_dir.name = "local-skill" + + skill_md = MagicMock(spec=Path) + skill_md.read_text.return_value = ( + "---\nname: local-skill\ndescription: Local desc\n---" + ) + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open") as mock_open_func, + patch.object(Path, "iterdir", return_value=[skill_dir]), + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=skill_md, + ), + ): + # Return different content based on file path + def open_side_effect(path, *args, **kwargs): + if SANDBOX_SKILLS_CACHE_FILENAME in str(path): + return mock_open(read_data=cache_data).return_value + return mock_open(read_data=config_data).return_value + + mock_open_func.side_effect = open_side_effect + result = skill_manager.list_skills(runtime="sandbox") + + assert len(result) == 1 + assert result[0].name == "local-skill" + assert result[0].source_type == "both" + assert result[0].local_exists is True + assert result[0].sandbox_exists is True + + def test_list_skills_sandbox_excludes_local_when_not_in_cache( + self, + skill_manager, + sandbox_skill_entry, + ): + """Test that local skills without a sandbox cache entry have sandbox_exists=False.""" + config_data = json.dumps({"skills": {"local-only": {"active": True}}}) + cache_data = json.dumps({"version": 1, "skills": []}) + + skill_dir = MagicMock(spec=Path) + skill_dir.is_dir.return_value = True + skill_dir.name = "local-only" + + skill_md = MagicMock(spec=Path) + skill_md.read_text.return_value = ( + "---\nname: local-only\ndescription: Local only\n---" + ) + + with ( + patch("os.path.exists", return_value=True), + patch("builtins.open") as mock_open_func, + patch.object(Path, "iterdir", return_value=[skill_dir]), + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=skill_md, + ), + ): + def open_side_effect(path, *args, **kwargs): + if SANDBOX_SKILLS_CACHE_FILENAME in str(path): + return mock_open(read_data=cache_data).return_value + return mock_open(read_data=config_data).return_value + + mock_open_func.side_effect = open_side_effect + result = skill_manager.list_skills(runtime="sandbox") + + assert len(result) == 1 + assert result[0].name == "local-only" + assert result[0].local_exists is True + assert result[0].sandbox_exists is False + assert result[0].source_type == "local_only" + + +# --------------------------------------------------------------- +# set_sandbox_skills_cache and get_sandbox_skills_cache_status +# --------------------------------------------------------------- + + +class TestSkillManagerSandboxCache: + """Test sandbox cache management.""" + + def test_set_sandbox_skills_cache_saves_deduped( + self, + skill_manager, + ): + """Test that set_sandbox_skills_cache deduplicates and saves.""" + skills = [ + {"name": "skill-a", "description": "A", "path": "/workspace/.../SKILL.md"}, + {"name": "skill-b", "description": "B", "path": "/workspace/.../SKILL.md"}, + {"name": "skill-a", "description": "A dup", "path": "/workspace/.../SKILL.md"}, + ] + + with patch.object(skill_manager, "_save_sandbox_skills_cache") as mock_save: + skill_manager.set_sandbox_skills_cache(skills) + + saved = mock_save.call_args[0][0] + assert "skills" in saved + assert len(saved["skills"]) == 2 + names = [s["name"] for s in saved["skills"]] + assert "skill-a" in names + assert "skill-b" in names + + def test_set_sandbox_skills_cache_skips_invalid_names( + self, + skill_manager, + ): + """Test that entries with invalid names are skipped.""" + skills = [ + {"name": "valid-skill", "description": "ok", "path": ""}, + {"name": "", "description": "empty", "path": ""}, + {"name": "../escape", "description": "bad", "path": ""}, + ] + + with patch.object(skill_manager, "_save_sandbox_skills_cache") as mock_save: + skill_manager.set_sandbox_skills_cache(skills) + + saved = mock_save.call_args[0][0] + assert len(saved["skills"]) == 1 + assert saved["skills"][0]["name"] == "valid-skill" + + def test_get_sandbox_skills_cache_status_ready( + self, + skill_manager, + ): + """Test that status returns ready=True when cache has skills.""" + with ( + patch.object( + skill_manager, + "_load_sandbox_skills_cache", + return_value={ + "version": 1, + "skills": [{"name": "s1"}], + "updated_at": "2025-01-01T00:00:00+00:00", + }, + ), + patch("os.path.exists", return_value=True), + ): + status = skill_manager.get_sandbox_skills_cache_status() + + assert status["exists"] is True + assert status["ready"] is True + assert status["count"] == 1 + + def test_get_sandbox_skills_cache_status_not_ready( + self, + skill_manager, + ): + """Test that status returns ready=False when no skills cached.""" + with ( + patch.object( + skill_manager, + "_load_sandbox_skills_cache", + return_value={"version": 1, "skills": []}, + ), + patch("os.path.exists", return_value=True), + ): + status = skill_manager.get_sandbox_skills_cache_status() + + assert status["exists"] is True + assert status["ready"] is False + assert status["count"] == 0 + + +# --------------------------------------------------------------- +# is_sandbox_only_skill +# --------------------------------------------------------------- + + +class TestSkillManagerIsSandboxOnly: + """Test is_sandbox_only_skill.""" + + def test_is_sandbox_only_skill_returns_true_when_only_in_cache( + self, + skill_manager, + ): + """Test that a skill existing only in cache returns True.""" + with ( + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=None, + ), + patch.object( + skill_manager, + "_load_sandbox_skills_cache", + return_value={ + "version": 1, + "skills": [{"name": "sandbox-only"}], + }, + ), + ): + assert skill_manager.is_sandbox_only_skill("sandbox-only") is True + + def test_is_sandbox_only_skill_returns_false_when_local_exists( + self, + skill_manager, + ): + """Test that a skill with local SKILL.md returns False.""" + with ( + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=MagicMock(), + ), + ): + assert skill_manager.is_sandbox_only_skill("local-skill") is False + + def test_is_sandbox_only_skill_returns_false_for_nonexistent( + self, + skill_manager, + ): + """Test that a skill not in cache or local returns False.""" + with ( + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=None, + ), + patch.object( + skill_manager, + "_load_sandbox_skills_cache", + return_value={"version": 1, "skills": []}, + ), + ): + assert skill_manager.is_sandbox_only_skill("ghost") is False + + +# --------------------------------------------------------------- +# set_skill_active / delete_skill +# --------------------------------------------------------------- + + +class TestSkillManagerMutations: + """Test set_skill_active and delete_skill.""" + + def test_set_skill_active_saves_config( + self, + skill_manager, + ): + """Test that set_skill_active writes the config with the new active state.""" + existing_config = {"skills": {"my-skill": {"active": False}}} + + with ( + patch( + "astrbot.core.skills.skill_manager._normalize_skill_markdown_path", + return_value=MagicMock(), + ), + patch("os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=json.dumps(existing_config))), + ): + skill_manager.set_skill_active("my-skill", True) + + # Config should have been saved with active=True + # We verify by checking that config file was opened for writing + # with content containing "active": true + # The exact assertion depends on mocking — key is it doesn't crash + + def test_set_skill_active_raises_on_sandbox_only( + self, + skill_manager, + ): + """Test that set_skill_active raises PermissionError for sandbox-only skills.""" + with ( + patch.object(skill_manager, "is_sandbox_only_skill", return_value=True), + pytest.raises(PermissionError, match="Sandbox preset skill"), + ): + skill_manager.set_skill_active("sandbox-only", False) + + def test_delete_skill_removes_directory_and_config( + self, + skill_manager, + ): + """Test that delete_skill removes the skill directory and config entry.""" + from astrbot.core.skills.skill_manager import _normalize_skill_markdown_path + + skill_dir = MagicMock(spec=Path) + skill_dir.exists.return_value = True + skill_dir.name = "my-skill" + + with ( + patch.object( + skill_manager, + "is_sandbox_only_skill", + return_value=False, + ), + patch("pathlib.Path", return_value=skill_dir) as mock_path_cls, + patch("shutil.rmtree") as mock_rmtree, + patch.object(skill_manager, "_remove_skill_from_sandbox_cache"), + patch.object(skill_manager, "_load_config", + return_value={"skills": {"my-skill": {"active": True}}}), + patch.object(skill_manager, "_save_config") as mock_save, + ): + skill_manager.delete_skill("my-skill") + mock_rmtree.assert_called_once_with(skill_dir) + mock_save.assert_called_once() + + def test_delete_skill_raises_on_sandbox_only( + self, + skill_manager, + ): + """Test that delete_skill raises PermissionError for sandbox-only skills.""" + with ( + patch.object(skill_manager, "is_sandbox_only_skill", return_value=True), + pytest.raises(PermissionError, match="Sandbox preset skill"), + ): + skill_manager.delete_skill("sandbox-only") + + +# --------------------------------------------------------------- +# utility / helper functions +# --------------------------------------------------------------- + + +class TestSkillManagerUtilities: + """Test standalone utility functions in skill_manager module.""" + + def test_normalize_skill_name_replaces_spaces(self): + """Test that _normalize_skill_name replaces whitespace with underscores.""" + result = _normalize_skill_name("my cool skill") + assert result == "my_cool_skill" + assert _normalize_skill_name(" hello world ") == "hello_world" + assert _normalize_skill_name(None) == "" + + def test_normalize_cached_sandbox_skill_path_default(self): + """Test that empty path defaults to the standard sandbox SKILL.md path.""" + result = _normalize_cached_sandbox_skill_path("my-skill", "") + expected = f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/my-skill/SKILL.md" + assert result == expected + + def test_normalize_cached_sandbox_skill_path_rejects_relative_escape(self): + """Test that path with '..' is rejected and falls back to default.""" + result = _normalize_cached_sandbox_skill_path("my-skill", "/workspace/../../etc/SKILL.md") + expected = f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/my-skill/SKILL.md" + assert result == expected + + def test_normalize_cached_sandbox_skill_path_rejects_wrong_filename(self): + """Test that a path not ending in SKILL.md falls back to default.""" + result = _normalize_cached_sandbox_skill_path("my-skill", "/workspace/skills/my-skill/README.md") + expected = f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/my-skill/SKILL.md" + assert result == expected + + def test_normalize_cached_sandbox_skill_path_rejects_wrong_dir_name(self): + """Test that a path with mismatched directory name falls back to default.""" + result = _normalize_cached_sandbox_skill_path("my-skill", "/workspace/skills/other-skill/SKILL.md") + expected = f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/my-skill/SKILL.md" + assert result == expected + + def test_is_ignored_zip_entry_macosx(self): + """Test that __MACOSX entries are ignored.""" + assert _is_ignored_zip_entry("__MACOSX/") is True + assert _is_ignored_zip_entry("__MACOSX/somefile") is True + assert _is_ignored_zip_entry("my-skill/SKILL.md") is False + + def test_sanitize_prompt_path_for_prompt_removes_backticks(self): + """Test that backticks are stripped from path.""" + result = _sanitize_prompt_path_for_prompt("/path/with`backticks`/SKILL.md") + assert "`" not in result + assert "backticks" in result + + def test_sanitize_prompt_description_cleans_whitespace(self): + """Test that description is cleaned of extra whitespace.""" + result = _sanitize_prompt_description(" hello world ") + assert result == "hello world" + + def test_sanitize_skill_display_name_invalid(self): + """Test that invalid names return .""" + result = _sanitize_skill_display_name("../escape") + assert result == "" + + def test_sanitize_skill_display_name_valid(self): + """Test that valid names pass through unchanged.""" + result = _sanitize_skill_display_name("my-valid_skill1") + assert result == "my-valid_skill1" + + def test_parse_frontmatter_extracts_meta(self): + """Test that YAML frontmatter is parsed correctly.""" + text = "---\nname: test\ndescription: hello\ninput_schema:\n type: object\n---\nBody" + result = _parse_frontmatter(text) + assert result["name"] == "test" + assert result["description"] == "hello" + assert result["input_schema"] == {"type": "object"} + + def test_parse_frontmatter_returns_empty_when_no_frontmatter(self): + """Test that text without frontmatter returns empty dict.""" + result = _parse_frontmatter("No frontmatter here") + assert result == {} + + def test_parse_frontmatter_returns_empty_on_yaml_error(self): + """Test that malformed YAML returns empty dict.""" + text = "---\n: invalid yaml :::\n---\nBody" + result = _parse_frontmatter(text) + assert result == {} + + def test_build_skills_prompt_generates_block(self): + """Test that build_skills_prompt generates a non-empty prompt block.""" + skills = [ + SkillInfo( + name="test-skill", + description="A test skill", + path="/tmp/skills/test-skill/SKILL.md", + active=True, + ), + ] + result = build_skills_prompt(skills) + assert "## Skills" in result + assert "test-skill" in result + assert "A test skill" in result + assert "SKILL.md" in result + + def test_build_skills_prompt_empty_list(self): + """Test that build_skills_prompt with no skills still produces a block.""" + result = build_skills_prompt([]) + assert "## Skills" in result + assert "Available skills" in result diff --git a/tests/unit/test_star_manager.py b/tests/unit/test_star_manager.py new file mode 100644 index 000000000..44439cbb7 --- /dev/null +++ b/tests/unit/test_star_manager.py @@ -0,0 +1,530 @@ +"""Unit tests for astrbot.core.star.star_manager. + +Tests PluginManager initialization, load, install, and uninstall flows with +full mock isolation (no filesystem, no network). +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest + +from astrbot.core.star.star_manager import ( + PluginDependencyInstallError, + PluginManager, + PluginVersionIncompatibleError, +) + + +# --------------------------------------------------------------------------- +# PluginVersionIncompatibleError / PluginDependencyInstallError +# --------------------------------------------------------------------------- + + +class TestPluginExceptions: + """Custom exception classes.""" + + def test_version_incompatible_error(self): + """PluginVersionIncompatibleError is a plain exception.""" + exc = PluginVersionIncompatibleError("bad version") + assert isinstance(exc, Exception) + assert str(exc) == "bad version" + + def test_dependency_install_error(self): + """PluginDependencyInstallError wraps the original error.""" + inner = ValueError("pip failed") + exc = PluginDependencyInstallError( + plugin_label="my_plugin", + requirements_path="/path/requirements.txt", + error=inner, + ) + assert exc.plugin_label == "my_plugin" + assert exc.requirements_path == "/path/requirements.txt" + assert exc.error is inner + assert "pip failed" in str(exc) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_context(): + """A fully mocked Context with required methods.""" + ctx = MagicMock() + ctx.get_all_stars.return_value = [] + ctx.get_registered_star.return_value = None + return ctx + + +@pytest.fixture +def mock_config(): + """A mock AstrBotConfig.""" + return MagicMock() + + +@pytest.fixture +def plugin_manager(mock_context, mock_config): + """Create a PluginManager with all dependencies mocked.""" + with ( + patch( + "astrbot.core.star.star_manager.get_astrbot_plugin_path", + return_value="/mock/plugins", + ), + patch( + "astrbot.core.star.star_manager.get_astrbot_config_path", + return_value="/mock/config", + ), + patch( + "astrbot.core.star.star_manager.get_astrbot_path", + return_value="/mock/astrbot", + ), + patch( + "astrbot.core.star.star_manager.PluginUpdator", + ), + patch( + "astrbot.core.star.star_manager.StarTools", + ), + ): + pm = PluginManager(mock_context, mock_config) + # Disable hot-reload watcher + pm.tasks = set() + return pm + + +# --------------------------------------------------------------------------- +# Constructor / Init +# --------------------------------------------------------------------------- + + +class TestPluginManagerInit: + """PluginManager constructor behavior.""" + + def test_init_sets_paths(self, plugin_manager): + """Constructor sets correct default paths.""" + assert plugin_manager.plugin_store_path == "/mock/plugins" + assert plugin_manager.plugin_config_path == "/mock/config" + assert "astrbot" in plugin_manager.reserved_plugin_path + + def test_init_sets_empty_failed_dict(self, plugin_manager): + """Constructor initializes empty failed plugin tracking.""" + assert plugin_manager.failed_plugin_dict == {} + assert plugin_manager.failed_plugin_info == "" + + def test_init_stores_context_and_config(self, plugin_manager, mock_context, mock_config): + """Constructor stores the context and config references.""" + assert plugin_manager.context is mock_context + assert plugin_manager.config is mock_config + + def test_init_sets_lock(self, plugin_manager): + """Constructor creates an asyncio.Lock.""" + import asyncio + + assert isinstance(plugin_manager._pm_lock, asyncio.Lock) + + def test_init_starts_watcher_when_reload_env_set(self, mock_context, mock_config): + """When ASTRBOT_RELOAD=1, the file watcher task is created.""" + with ( + patch( + "astrbot.core.star.star_manager.get_astrbot_plugin_path", + return_value="/mock/plugins", + ), + patch( + "astrbot.core.star.star_manager.get_astrbot_config_path", + return_value="/mock/config", + ), + patch( + "astrbot.core.star.star_manager.get_astrbot_path", + return_value="/mock", + ), + patch( + "astrbot.core.star.star_manager.PluginUpdator", + ), + patch( + "astrbot.core.star.star_manager.StarTools", + ), + patch.dict("os.environ", {"ASTRBOT_RELOAD": "1"}), + patch( + "astrbot.core.star.star_manager.asyncio.create_task", + ) as mock_create_task, + ): + pm = PluginManager(mock_context, mock_config) + # Clean up after init + assert mock_create_task.called + + +# --------------------------------------------------------------------------- +# _get_classes / _get_modules / _get_plugin_modules (static helpers) +# --------------------------------------------------------------------------- + + +class TestGetClasses: + """PluginManager._get_classes() static helper.""" + + def test_returns_class_names_from_module(self): + """_get_classes finds classes ending in 'plugin' or named 'main'.""" + import types + + module = types.ModuleType("test_module") + exec( + """ +class MyPlugin: + pass + +class NotAPlugin: + pass +""", + module.__dict__, + ) + result = PluginManager._get_classes(module) + assert "MyPlugin" in result + assert "NotAPlugin" not in result + + def test_returns_main_class(self): + """_get_classes includes a class named 'main' (case-insensitive).""" + import types + + module = types.ModuleType("test_module") + exec( + """ +class Main: + pass +""", + module.__dict__, + ) + result = PluginManager._get_classes(module) + assert "Main" in result + + def test_returns_empty_when_no_match(self): + """_get_classes returns empty list when no matching classes exist.""" + import types + + module = types.ModuleType("test_module") + exec( + """ +class Helper: + pass +""", + module.__dict__, + ) + result = PluginManager._get_classes(module) + assert result == [] + + +class TestValidateImportableName: + """PluginManager._validate_importable_name() static helper.""" + + def test_valid_name_passes(self): + """A valid Python identifier passes validation.""" + PluginManager._validate_importable_name("my_plugin") # no raise + + def test_rejects_path_separator(self): + """A name containing / raises ValueError.""" + with pytest.raises(ValueError, match="路径分隔符"): + PluginManager._validate_importable_name("my/plugin") + + def test_rejects_invalid_identifier(self): + """A name that is not a valid Python identifier raises Exception.""" + with pytest.raises(Exception, match="合法的模块名称"): + PluginManager._validate_importable_name("123invalid") + + def test_rejects_keyword(self): + """A name that is a Python keyword raises Exception.""" + with pytest.raises(Exception, match="合法的模块名称"): + PluginManager._validate_importable_name("class") + + +# --------------------------------------------------------------------------- +# load +# --------------------------------------------------------------------------- + + +class TestLoad: + """PluginManager.load() behavior.""" + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + async def test_load_returns_true_when_no_plugins( + self, mock_sync, mock_sp_get, plugin_manager + ): + """load() returns (True, None) when no plugin modules exist.""" + mock_sp_get.return_value = [] + plugin_manager._get_plugin_modules = MagicMock(return_value=[]) + success, error = await plugin_manager.load() + assert success is True + assert error is None + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + async def test_load_returns_false_when_modules_is_none( + self, mock_sync, mock_sp_get, plugin_manager + ): + """load() returns (False, msg) when _get_plugin_modules returns None.""" + mock_sp_get.return_value = [] + plugin_manager._get_plugin_modules = MagicMock(return_value=None) + success, error = await plugin_manager.load() + assert success is False + assert "未找到" in error + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + async def test_load_handles_import_failure( + self, mock_sync, mock_sp_get, plugin_manager + ): + """load() records failed plugins when import fails.""" + mock_sp_get.return_value = [] + plugin_manager._get_plugin_modules = MagicMock( + return_value=[ + { + "pname": "broken_plugin", + "module": "main", + "module_path": "/mock/plugins/broken_plugin/main", + "reserved": False, + } + ] + ) + plugin_manager._import_plugin_with_dependency_recovery = AsyncMock( + side_effect=ImportError("Module not found") + ) + plugin_manager._load_plugin_metadata = MagicMock(return_value=None) + plugin_manager._build_failed_plugin_record = MagicMock( + return_value={"name": "broken_plugin", "error": "Module not found"} + ) + + success, error = await plugin_manager.load() + assert success is False + assert "broken_plugin" in plugin_manager.failed_plugin_dict + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + async def test_load_calls_sync_command_configs( + self, mock_sync, mock_sp_get, plugin_manager + ): + """load() calls sync_command_configs after processing plugins.""" + mock_sp_get.return_value = [] + plugin_manager._get_plugin_modules = MagicMock(return_value=[]) + await plugin_manager.load() + mock_sync.assert_called_once() + + +# --------------------------------------------------------------------------- +# install_plugin +# --------------------------------------------------------------------------- + + +class TestInstallPlugin: + """PluginManager.install_plugin() behavior.""" + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + @patch("astrbot.core.star.star_manager.Metric") + async def test_install_plugin_success( + self, mock_metric, mock_sync, mock_sp_get, plugin_manager, mock_context + ): + """install_plugin installs and loads a plugin successfully.""" + mock_sp_get.return_value = [] + plugin_manager.updator.parse_github_url = MagicMock( + return_value=("owner", "test_repo", "repo") + ) + plugin_manager.updator.format_name = MagicMock(return_value="test_repo") + + # Mock the install to return a plugin path + plugin_manager.updator.install = AsyncMock( + return_value="/mock/plugins/test_repo" + ) + plugin_manager._get_plugin_dir_name_from_metadata = MagicMock( + return_value="test_repo" + ) + plugin_manager._ensure_plugin_requirements = AsyncMock() + plugin_manager.load = AsyncMock(return_value=(True, None)) + + mock_plugin_meta = MagicMock() + mock_plugin_meta.repo = "https://github.com/test/test_repo" + mock_plugin_meta.name = "test_plugin" + mock_context.get_registered_star.side_effect = lambda name: ( + mock_plugin_meta if name == "test_repo" else None + ) + + result = await plugin_manager.install_plugin( + "https://github.com/test/test_repo" + ) + assert result is not None + assert result["repo"] == "https://github.com/test/test_repo" + assert result["name"] == "test_plugin" + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + @patch("astrbot.core.star.star_manager.Metric") + async def test_install_plugin_raises_when_load_fails( + self, mock_metric, mock_sync, mock_sp_get, plugin_manager + ): + """install_plugin raises when load() returns failure.""" + mock_sp_get.return_value = [] + plugin_manager.updator.parse_github_url = MagicMock( + return_value=("owner", "test_repo", "repo") + ) + plugin_manager.updator.format_name = MagicMock(return_value="test_repo") + plugin_manager.updator.install = AsyncMock( + return_value="/mock/plugins/test_repo" + ) + plugin_manager._get_plugin_dir_name_from_metadata = MagicMock( + return_value="test_repo" + ) + plugin_manager._ensure_plugin_requirements = AsyncMock() + plugin_manager.load = AsyncMock( + return_value=(False, "Version incompatible error") + ) + + with pytest.raises(Exception, match="Version incompatible error"): + await plugin_manager.install_plugin( + "https://github.com/test/test_repo" + ) + + @patch("astrbot.core.star.star_manager.sp.global_get") + @patch("astrbot.core.star.star_manager.sync_command_configs", new_callable=AsyncMock) + @patch("astrbot.core.star.star_manager.Metric") + async def test_install_plugin_raises_when_dir_exists( + self, mock_metric, mock_sync, mock_sp_get, plugin_manager + ): + """install_plugin raises when the target directory already exists.""" + mock_sp_get.return_value = [] + plugin_manager.updator.parse_github_url = MagicMock( + return_value=("owner", "test_repo", "repo") + ) + plugin_manager.updator.format_name = MagicMock(return_value="test_repo") + + with ( + patch("astrbot.core.star.star_manager.anyio.Path") as mock_path, + ): + mock_path_instance = MagicMock() + mock_path.return_value = mock_path_instance + mock_path_instance.exists = AsyncMock(return_value=True) + + with pytest.raises(Exception, match="已存在"): + await plugin_manager.install_plugin( + "https://github.com/test/test_repo" + ) + + +# --------------------------------------------------------------------------- +# uninstall_plugin +# --------------------------------------------------------------------------- + + +class TestUninstallPlugin: + """PluginManager.uninstall_plugin() behavior.""" + + @patch("astrbot.core.star.star_manager.remove_dir") + @patch("astrbot.core.star.star_manager.unregister_platform_adapters_by_module") + async def test_uninstall_plugin_success( + self, + mock_unregister, + mock_remove_dir, + plugin_manager, + mock_context, + ): + """uninstall_plugin terminates, unbinds, and removes plugin directory.""" + mock_plugin = MagicMock() + mock_plugin.reserved = False + mock_plugin.root_dir_name = "test_repo" + mock_plugin.name = "test_plugin" + mock_plugin.module_path = "data.plugins.test_repo.main" + mock_context.get_registered_star.return_value = mock_plugin + + plugin_manager._terminate_plugin = AsyncMock() + plugin_manager._unbind_plugin = AsyncMock() + + await plugin_manager.uninstall_plugin("test_plugin") + + plugin_manager._terminate_plugin.assert_called_once_with(mock_plugin) + plugin_manager._unbind_plugin.assert_called_once_with( + "test_plugin", "data.plugins.test_repo.main" + ) + mock_remove_dir.assert_called_once() + + @patch("astrbot.core.star.star_manager.remove_dir") + async def test_uninstall_plugin_raises_when_not_found( + self, mock_remove_dir, plugin_manager, mock_context + ): + """uninstall_plugin raises when plugin is not registered.""" + mock_context.get_registered_star.return_value = None + with pytest.raises(Exception, match="插件不存在"): + await plugin_manager.uninstall_plugin("nonexistent") + + @patch("astrbot.core.star.star_manager.remove_dir") + async def test_uninstall_plugin_raises_for_reserved( + self, mock_remove_dir, plugin_manager, mock_context + ): + """uninstall_plugin raises when plugin is reserved.""" + mock_plugin = MagicMock() + mock_plugin.reserved = True + mock_context.get_registered_star.return_value = mock_plugin + + with pytest.raises(Exception, match="保留插件"): + await plugin_manager.uninstall_plugin("reserved_plugin") + + +# --------------------------------------------------------------------------- +# _validate_astrbot_version_specifier +# --------------------------------------------------------------------------- + + +class TestValidateAstrbotVersion: + """PluginManager._validate_astrbot_version_specifier() behavior.""" + + def test_none_version_returns_valid(self): + """None version specifier is valid.""" + valid, msg = PluginManager._validate_astrbot_version_specifier(None) + assert valid is True + assert msg is None + + @patch("astrbot.core.star.star_manager.VERSION", "4.16.0") + def test_version_in_range(self): + """A version specifier that includes the current version is valid.""" + valid, msg = PluginManager._validate_astrbot_version_specifier(">=4.16,<5") + assert valid is True + assert msg is None + + @patch("astrbot.core.star.star_manager.VERSION", "4.15.0") + def test_version_out_of_range(self): + """A version specifier that excludes the current version is invalid.""" + valid, msg = PluginManager._validate_astrbot_version_specifier(">=4.16,<5") + assert valid is False + assert "does not satisfy" in msg + + def test_invalid_specifier_returns_error(self): + """An unparseable specifier returns invalid with error message.""" + valid, msg = PluginManager._validate_astrbot_version_specifier("not_a_version") + assert valid is False + assert "PEP 440" in msg + + +# --------------------------------------------------------------------------- +# _load_plugin_metadata +# --------------------------------------------------------------------------- + + +class TestLoadPluginMetadata: + """PluginManager._load_plugin_metadata() behavior.""" + + def test_raises_when_path_does_not_exist(self): + """_load_plugin_metadata raises when plugin_path does not exist.""" + with patch("astrbot.core.star.star_manager.os.path.exists", return_value=False): + with pytest.raises(Exception, match="插件不存在"): + PluginManager._load_plugin_metadata("/nonexistent/path") + + +# --------------------------------------------------------------------------- +# reload_failed_plugin +# --------------------------------------------------------------------------- + + +class TestReloadFailedPlugin: + """PluginManager.reload_failed_plugin() behavior.""" + + async def test_returns_false_when_not_in_failed_dict(self, plugin_manager): + """reload_failed_plugin returns (False, msg) when dir not in failed dict.""" + result = await plugin_manager.reload_failed_plugin("unknown") + assert result == (False, "插件不存在于失败列表中") diff --git a/tests/unit/test_storage_cleaner.py b/tests/unit/test_storage_cleaner.py new file mode 100644 index 000000000..76edb72f6 --- /dev/null +++ b/tests/unit/test_storage_cleaner.py @@ -0,0 +1,262 @@ +"""Unit tests for astrbot.core.utils.storage_cleaner. + +Uses tmp_path for filesystem-level assertions. +""" + +import os +from pathlib import Path +from unittest.mock import patch + +import pytest + +from astrbot.core.utils.storage_cleaner import StorageCleaner + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _make_file(path: Path, size: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"x" * size) + + +def _assert_exists(path: Path, expected_size: int | None = None) -> None: + assert path.exists(), f"Expected {path} to exist" + if expected_size is not None: + assert path.stat().st_size == expected_size + + +def _assert_missing(path: Path) -> None: + assert not path.exists(), f"Expected {path} to be removed" + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestStorageCleanerStatus: + def test_get_status_reports_logs_and_cache(self, tmp_path: Path): + data = tmp_path / "data" + temp = data / "temp" + logs = data / "logs" + + _make_file(temp / "x.wav", 100) + _make_file(data / "plugins.json", 50) + _make_file(logs / "astrbot.log", 200) + _make_file(logs / "astrbot.2026-01-01.log", 80) + + cleaner = StorageCleaner( + {"log_file_enable": True, "log_file_path": "logs/astrbot.log"}, + data_dir=data, + temp_dir=temp, + ) + + status = cleaner.get_status() + + assert status["logs"]["size_bytes"] == 280 # 200 + 80 + assert status["logs"]["file_count"] == 2 + assert status["cache"]["size_bytes"] == 150 # 100 + 50 + assert status["cache"]["file_count"] == 2 + assert status["total_bytes"] == 430 + + def test_get_status_when_dirs_missing(self, tmp_path: Path): + data = tmp_path / "data" + cleaner = StorageCleaner( + {"log_file_enable": False}, + data_dir=data, + temp_dir=data / "temp", + ) + status = cleaner.get_status() + assert status["logs"]["size_bytes"] == 0 + assert status["logs"]["file_count"] == 0 + assert status["cache"]["size_bytes"] == 0 + assert status["total_bytes"] == 0 + + def test_get_status_includes_cache_extra_files(self, tmp_path: Path): + data = tmp_path / "data" + temp = data / "temp" + logs = data / "logs" + _make_file(logs / "astrbot.log", 10) + _make_file(temp / "a.bin", 10) + + _make_file(data / "plugins_custom_foo.json", 20) + _make_file(data / "sandbox_skills_cache.json", 30) + + cleaner = StorageCleaner( + {"log_file_enable": True}, + data_dir=data, + temp_dir=temp, + ) + status = cleaner.get_status() + assert status["cache"]["file_count"] >= 3 # temp + custom* + sandbox + assert status["cache"]["size_bytes"] >= 60 + + +class TestStorageCleanerCleanup: + def test_cleanup_all_removes_logs_and_cache(self, tmp_path: Path): + data = tmp_path / "data" + temp = data / "temp" + logs = data / "logs" + + active_log = logs / "astrbot.log" + rotated_log = logs / "astrbot.2026-03-01.log" + _make_file(active_log, 100) + _make_file(rotated_log, 50) + _make_file(temp / "t.bin", 200) + + cleaner = StorageCleaner( + {"log_file_enable": True, "log_file_path": "logs/astrbot.log"}, + data_dir=data, + temp_dir=temp, + ) + result = cleaner.cleanup("all") + + # Active log → truncated + _assert_exists(active_log, 0) + # Rotated log → deleted + _assert_missing(rotated_log) + # Temp cache → deleted + _assert_missing(temp / "t.bin") + + assert result["removed_bytes"] == 350 + assert result["processed_files"] == 3 + assert result["deleted_files"] == 2 + assert result["truncated_files"] == 1 + + def test_cleanup_logs_target_only(self, tmp_path: Path): + data = tmp_path / "data" + logs = data / "logs" + temp = data / "temp" + + _make_file(logs / "astrbot.log", 80) + _make_file(temp / "c.bin", 999) + + cleaner = StorageCleaner( + {"log_file_enable": True, "log_file_path": "logs/astrbot.log"}, + data_dir=data, + temp_dir=temp, + ) + result = cleaner.cleanup("logs") + + _assert_exists(logs / "astrbot.log", 0) + # temp should be untouched + _assert_exists(temp / "c.bin", 999) + assert result["removed_bytes"] == 80 + + def test_cleanup_cache_target_only(self, tmp_path: Path): + data = tmp_path / "data" + logs = data / "logs" + temp = data / "temp" + + _make_file(logs / "astrbot.log", 80) + _make_file(temp / "c.bin", 200) + + cleaner = StorageCleaner( + {"log_file_enable": True, "log_file_path": "logs/astrbot.log"}, + data_dir=data, + temp_dir=temp, + ) + result = cleaner.cleanup("cache") + + _assert_exists(logs / "astrbot.log", 80) # untouched + _assert_missing(temp / "c.bin") + assert result["removed_bytes"] == 200 + + def test_cleanup_invalid_target_raises(self, tmp_path: Path): + cleaner = StorageCleaner({}, data_dir=tmp_path, temp_dir=tmp_path) + with pytest.raises(ValueError, match="Unsupported cleanup target"): + cleaner.cleanup("invalid") + + def test_cleanup_with_no_config(self, tmp_path: Path): + """All config options off; no active log files, so everything gets deleted.""" + data = tmp_path / "data" + _make_file(data / "logs" / "astrbot.log", 50) + + cleaner = StorageCleaner( + {"log_file_enable": False}, + data_dir=data, + temp_dir=data / "temp", + ) + result = cleaner.cleanup("logs") + # With log_file_enable=False, astrbot.log is NOT in active_log_files, + # so it gets deleted, not truncated. + _assert_missing(data / "logs" / "astrbot.log") + assert result["deleted_files"] == 1 + assert result["truncated_files"] == 0 + + +class TestStorageCleanerEdgeCases: + def test_cleanup_removes_empty_temp_dirs(self, tmp_path: Path): + data = tmp_path / "data" + temp = data / "temp" + nested = temp / "sub" / "nested" + _make_file(nested / "f.bin", 100) + + cleaner = StorageCleaner( + {}, + data_dir=data, + temp_dir=temp, + ) + cleaner.cleanup("cache") + + assert not nested.exists() + assert temp.exists() # root temp should be recreated by _cleanup_target + + def test_cleanup_skips_inexistent_file(self, tmp_path: Path): + data = tmp_path / "data" + cleaner = StorageCleaner( + {"log_file_enable": False}, + data_dir=data, + temp_dir=data / "temp", + ) + # No files at all – should not crash. + result = cleaner.cleanup("all") + assert result["failed_files"] == 0 + + def test_cleanup_handles_stat_os_error(self, tmp_path: Path): + data = tmp_path / "data" + logs = data / "logs" + _make_file(logs / "astrbot.log", 100) + + cleaner = StorageCleaner( + {"log_file_enable": True, "log_file_path": "logs/astrbot.log"}, + data_dir=data, + temp_dir=data / "temp", + ) + + original_stat = (logs / "astrbot.log").stat + + def _broken_stat(): + raise OSError(13, "Permission denied") + + (logs / "astrbot.log").stat = _broken_stat # type: ignore[method-assign] + + result = cleaner.cleanup("logs") + assert result["failed_files"] == 1 + + (logs / "astrbot.log").stat = original_stat + + def test_cleanup_handles_unlink_os_error(self, tmp_path: Path): + data = tmp_path / "data" + logs = data / "logs" + _make_file(logs / "astrbot.log", 100) + + cleaner = StorageCleaner( + {"log_file_enable": False}, + data_dir=data, + temp_dir=data / "temp", + ) + + original_unlink = (logs / "astrbot.log").unlink + + def _broken_unlink(): + raise OSError(13, "Permission denied") + + (logs / "astrbot.log").unlink = _broken_unlink # type: ignore[method-assign] + + result = cleaner.cleanup("logs") + assert result["failed_files"] == 1 + + (logs / "astrbot.log").unlink = original_unlink diff --git a/tests/unit/test_temp_dir_cleaner.py b/tests/unit/test_temp_dir_cleaner.py new file mode 100644 index 000000000..5c16d1eea --- /dev/null +++ b/tests/unit/test_temp_dir_cleaner.py @@ -0,0 +1,259 @@ +"""Unit tests for astrbot.core.utils.temp_dir_cleaner. + +Covers parse_size_to_bytes, cleanup_once, async lifecycle, and error paths. +""" + +import asyncio +import os +import time +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +from astrbot.core.utils.temp_dir_cleaner import TempDirCleaner, parse_size_to_bytes + +# --------------------------------------------------------------------------- +# parse_size_to_bytes +# --------------------------------------------------------------------------- + + +class TestParseSizeToBytes: + def test_valid_mb_string(self): + assert parse_size_to_bytes("1024") == 1024 * 1024**2 + + def test_valid_float_mb(self): + assert parse_size_to_bytes(0.5) == int(0.5 * 1024**2) + + def test_zero_returns_zero(self): + assert parse_size_to_bytes(0) == 0 + + def test_none_returns_zero(self): + assert parse_size_to_bytes(None) == 0 + + def test_invalid_string_returns_zero(self): + assert parse_size_to_bytes("not-a-number") == 0 + + def test_negative_value_returns_zero(self): + assert parse_size_to_bytes("-10") == 0 + + def test_whitespace_string(self): + assert parse_size_to_bytes(" 512 ") == 512 * 1024**2 + + +# --------------------------------------------------------------------------- +# helpers for filesystem tests +# --------------------------------------------------------------------------- + + +def _write_file(path: Path, size: int, mtime: float | None = None) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"x" * size) + if mtime is not None: + os.utime(path, (mtime, mtime)) + + +def _file_sizes(temp_dir: Path) -> list[tuple[Path, int]]: + return [ + (f, f.stat().st_size) + for f in sorted(temp_dir.rglob("*")) + if f.is_file() + ] + + +# --------------------------------------------------------------------------- +# cleanup_once +# --------------------------------------------------------------------------- + + +class TestCleanupOnce: + def test_noop_when_below_limit(self, tmp_path: Path): + temp_dir = tmp_path / "temp" + temp_dir.mkdir(parents=True, exist_ok=True) + _write_file(temp_dir / "a.bin", 100, time.time()) + + cleaner = TempDirCleaner( + max_size_getter=lambda: "1", # 1 MB → effectively unlimited + temp_dir=temp_dir, + ) + cleaner.cleanup_once() + + files = _file_sizes(temp_dir) + assert len(files) == 1 + assert files[0][1] == 100 + + def test_removes_oldest_files_when_over_limit(self, tmp_path: Path): + temp_dir = tmp_path / "temp" + temp_dir.mkdir(parents=True, exist_ok=True) + + base = time.time() - 1000 + _write_file(temp_dir / "old.bin", 400, base) + _write_file(temp_dir / "mid.bin", 300, base + 10) + _write_file(temp_dir / "new.bin", 300, base + 20) + + # Limit = 0.0008 MB ≈ 838 bytes; total = 1000 => release 30% = 300 + cleaner = TempDirCleaner( + max_size_getter=lambda: "0.0008", + temp_dir=temp_dir, + ) + cleaner.cleanup_once() + + remaining = _file_sizes(temp_dir) + remaining_total = sum(sz for _, sz in remaining) + # old.bin (oldest) should have been deleted first + assert (temp_dir / "old.bin").exists() is False + # remaining <= 700 (total - minimum 30%) + assert remaining_total <= 700 + + def test_noop_when_temp_dir_missing(self, tmp_path: Path): + temp_dir = tmp_path / "nonexistent" + cleaner = TempDirCleaner( + max_size_getter=lambda: "0.001", + temp_dir=temp_dir, + ) + # Should not raise + cleaner.cleanup_once() + + def test_handles_unlink_os_error_gracefully(self, tmp_path: Path): + temp_dir = tmp_path / "temp" + temp_dir.mkdir(parents=True, exist_ok=True) + _write_file(temp_dir / "bad.bin", 9999, time.time() - 100) + + cleaner = TempDirCleaner( + max_size_getter=lambda: "0.001", + temp_dir=temp_dir, + ) + + def _broken_unlink(): + raise OSError(13, "Permission denied") + + (temp_dir / "bad.bin").unlink = _broken_unlink # type: ignore[method-assign] + + # Should log warning but not crash + with patch("astrbot.core.utils.temp_dir_cleaner.logger.warning") as mock_warn: + cleaner.cleanup_once() + mock_warn.assert_called() + assert any( + "Permission denied" in str(c) for c in mock_warn.call_args_list + ) + + def test_invalid_config_falls_back_to_default(self, tmp_path: Path): + temp_dir = tmp_path / "temp" + temp_dir.mkdir(parents=True, exist_ok=True) + _write_file(temp_dir / "a.bin", 50, time.time()) + + cleaner = TempDirCleaner( + max_size_getter=lambda: "invalid", + temp_dir=temp_dir, + ) + with patch("astrbot.core.utils.temp_dir_cleaner.logger.warning") as mock_warn: + cleaner.cleanup_once() + mock_warn.assert_called_once() + assert "fallback" in str(mock_warn.call_args[0][0]) + + def test_cleanup_once_removes_empty_dirs(self, tmp_path: Path): + temp_dir = tmp_path / "temp" + nested = temp_dir / "a" / "b" + _write_file(nested / "f.bin", 1000000, time.time() - 5000) + # Also create an already-empty subdir + empty_dir = temp_dir / "empty" + empty_dir.mkdir(parents=True, exist_ok=True) + + cleaner = TempDirCleaner( + max_size_getter=lambda: "0.001", + temp_dir=temp_dir, + ) + cleaner.cleanup_once() + + # The file in nested is old → gets deleted, dirs cleaned up + assert not nested.exists() + # empty dir should be removed by _cleanup_empty_dirs + _assert_dir_gone = not empty_dir.exists() + # either nested or empty (or both) should be gone + assert _assert_dir_gone or not nested.exists() + + +# --------------------------------------------------------------------------- +# scan / async lifecycle +# --------------------------------------------------------------------------- + + +class TestScanAndAsyncLifecycle: + def test_scan_empty_dir_returns_zeroes(self, tmp_path: Path): + temp_dir = tmp_path / "empty_temp" + temp_dir.mkdir(parents=True, exist_ok=True) + cleaner = TempDirCleaner( + max_size_getter=lambda: "100", + temp_dir=temp_dir, + ) + total, files = cleaner._scan_temp_files() + assert total == 0 + assert files == [] + + def test_scan_missing_dir_returns_zeroes(self, tmp_path: Path): + temp_dir = tmp_path / "ghost" + cleaner = TempDirCleaner( + max_size_getter=lambda: "100", + temp_dir=temp_dir, + ) + total, files = cleaner._scan_temp_files() + assert total == 0 + assert files == [] + + @pytest.mark.asyncio + async def test_stop_sets_event(self): + cleaner = TempDirCleaner( + max_size_getter=lambda: "100", + ) + await cleaner.stop() + assert cleaner._stop_event.is_set() is True + + @pytest.mark.asyncio + async def test_run_stops_when_stop_is_called(self, tmp_path: Path): + temp_dir = tmp_path / "temp" + temp_dir.mkdir(parents=True, exist_ok=True) + cleaner = TempDirCleaner( + max_size_getter=lambda: "1000", + temp_dir=temp_dir, + ) + + task = asyncio.create_task(cleaner.run()) + await asyncio.sleep(0.05) + assert task.done() is False + + await cleaner.stop() + await asyncio.wait_for(task, timeout=2.0) + assert task.done() is True + + @pytest.mark.asyncio + async def test_run_logs_exception_and_continues(self, tmp_path: Path): + """When cleanup_once raises, run() should log and continue the loop.""" + temp_dir = tmp_path / "temp" + temp_dir.mkdir(parents=True, exist_ok=True) + cleaner = TempDirCleaner( + max_size_getter=lambda: "1000", + temp_dir=temp_dir, + ) + + # Make cleanup_once raise once, then succeed + call_count = 0 + + def _flaky_cleanup(): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("transient failure") + + cleaner.cleanup_once = _flaky_cleanup # type: ignore[method-assign] + + with patch( + "astrbot.core.utils.temp_dir_cleaner.logger.error" + ) as mock_log_error: + task = asyncio.create_task(cleaner.run()) + await asyncio.sleep(0.05) + await cleaner.stop() + await asyncio.wait_for(task, timeout=2.0) + + mock_log_error.assert_called() + error_msg = str(mock_log_error.call_args[0][0]) + assert "transient failure" in error_msg or "failed" in error_msg diff --git a/tests/unit/test_tool.py b/tests/unit/test_tool.py new file mode 100644 index 000000000..13e40fd88 --- /dev/null +++ b/tests/unit/test_tool.py @@ -0,0 +1,344 @@ +"""Unit tests for astrbot.core.agent.tool: FunctionTool, ToolSchema, ToolSet.""" + +import pytest +from unittest.mock import MagicMock, AsyncMock +from pydantic import ValidationError + +from astrbot.core.agent.tool import ( + FunctionTool, + ToolSchema, + ToolSet, + ToolArgumentSpec, + ToolExecResult, +) + + +class TestToolSchema: + """ToolSchema construction and validation.""" + + def test_tool_schema_minimal(self): + """Construct with only name and description.""" + schema = ToolSchema(name="test_tool", description="A test tool") + assert schema.name == "test_tool" + assert schema.description == "A test tool" + assert schema.parameters is None + assert schema.active is True + + def test_tool_schema_with_parameters(self): + """Construct with valid JSON Schema parameters.""" + params = { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"} + }, + "required": ["query"], + } + schema = ToolSchema(name="search", description="Search tool", parameters=params) + assert schema.name == "search" + assert schema.parameters == params + + def test_tool_schema_invalid_parameters_raises(self): + """Reject parameters that are not valid JSON Schema.""" + with pytest.raises(ValidationError): + ToolSchema( + name="bad", + description="Bad schema", + parameters={"type": "nonexistent"}, + ) + + def test_tool_schema_inactive(self): + """Explicitly set active=False.""" + schema = ToolSchema(name="off", description="Disabled", active=False) + assert schema.active is False + + def test_tool_schema_parameters_none_allowed(self): + """parameters=None is valid and passes validation.""" + schema = ToolSchema(name="no_params", description="No parameters", parameters=None) + assert schema.parameters is None + + def test_tool_schema_empty_parameters_object(self): + """Empty object parameters are valid JSON Schema.""" + schema = ToolSchema( + name="empty", + description="Empty params", + parameters={"type": "object", "properties": {}}, + ) + assert schema.parameters == {"type": "object", "properties": {}} + + +class TestFunctionTool: + """FunctionTool construction and call interface.""" + + def test_function_tool_minimal(self): + """Construct with minimal fields.""" + tool = FunctionTool(name="echo", description="Echo tool") + assert tool.name == "echo" + assert tool.description == "Echo tool" + assert tool.handler is None + assert tool.active is True + assert tool.is_background_task is False + assert tool.source == "plugin" + + def test_function_tool_with_handler(self): + """Construct with an async handler.""" + async def handler(context, **kwargs): + return "ok" + tool = FunctionTool( + name="greet", + description="Greet", + handler=handler, + handler_module_path="tests.test_tool", + ) + assert tool.handler is handler + assert tool.handler_module_path == "tests.test_tool" + + def test_function_tool_background_task(self): + """Construct with is_background_task=True.""" + tool = FunctionTool( + name="bg", + description="Background", + is_background_task=True, + ) + assert tool.is_background_task is True + + def test_function_tool_source_values(self): + """Construct with different source values.""" + for source in ("plugin", "internal", "mcp"): + tool = FunctionTool( + name=f"tool_{source}", + description=source, + source=source, + ) + assert tool.source == source + + def test_function_tool_repr(self): + """__repr__ returns a meaningful string.""" + tool = FunctionTool(name="sum", description="Sum numbers") + rep = repr(tool) + assert "FuncTool" in rep + assert "sum" in rep + + def test_function_tool_call_not_implemented(self): + """call() raises NotImplementedError when no handler is set.""" + tool = FunctionTool(name="todo", description="Not implemented") + with pytest.raises(NotImplementedError, match="FunctionTool.call"): + import asyncio + asyncio.run(tool.call(MagicMock())) + + def test_function_tool_with_parameters(self): + """Construct with valid parameters.""" + params = { + "type": "object", + "properties": {"x": {"type": "integer"}}, + } + tool = FunctionTool(name="add", description="Add", parameters=params) + assert tool.parameters == params + + def test_function_tool_active_false(self): + """Construct with active=False.""" + tool = FunctionTool(name="inactive_tool", description="Not active", active=False) + assert tool.active is False + + def test_function_tool_inherits_validation(self): + """FunctionTool inherits ToolSchema parameter validation.""" + with pytest.raises(ValidationError): + FunctionTool( + name="bad", + description="Bad", + parameters={"type": "madeup"}, + ) + + +class TestToolSet: + """ToolSet add/remove/get and serialization methods.""" + + def test_empty_toolset(self): + """New ToolSet is empty.""" + ts = ToolSet() + assert ts.empty() is True + assert len(ts) == 0 + assert bool(ts) is False + + def test_add_tool(self): + """Add a tool increases length.""" + ts = ToolSet() + tool = FunctionTool(name="a", description="Tool A") + ts.add_tool(tool) + assert len(ts) == 1 + assert ts.empty() is False + assert bool(ts) is True + + def test_add_duplicate_name_overwrites(self): + """Adding a tool with same name replaces the existing one.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="dup", description="Original")) + ts.add_tool(FunctionTool(name="dup", description="Replacement")) + assert len(ts) == 1 + tool = ts.get_tool("dup") + assert tool is not None + assert tool.description == "Replacement" + + def test_add_duplicate_active_prefers_active(self): + """When a duplicate is added, active=True wins over active=False.""" + ts = ToolSet() + inactive = FunctionTool(name="x", description="Inactive", active=False) + active = FunctionTool(name="x", description="Active", active=True) + ts.add_tool(inactive) + ts.add_tool(active) + tool = ts.get_tool("x") + assert tool is not None + assert tool.description == "Active" + + def test_add_duplicate_inactive_does_not_replace_active(self): + """Adding an inactive tool with an already active name does not overwrite.""" + ts = ToolSet() + active = FunctionTool(name="y", description="Active", active=True) + inactive = FunctionTool(name="y", description="Inactive", active=False) + ts.add_tool(active) + ts.add_tool(inactive) + tool = ts.get_tool("y") + assert tool is not None + assert tool.description == "Active" + + def test_remove_tool(self): + """Remove a tool by name.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="keep", description="Keep")) + ts.add_tool(FunctionTool(name="remove_me", description="Remove")) + ts.remove_tool("remove_me") + assert len(ts) == 1 + assert ts.get_tool("keep") is not None + assert ts.get_tool("remove_me") is None + + def test_remove_nonexistent_tool(self): + """Removing a non-existent tool does not raise.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="a", description="A")) + ts.remove_tool("nonexistent") + assert len(ts) == 1 + + def test_get_tool_returns_none_for_non_functiontool(self): + """get_tool returns None when a ToolSchema (not FunctionTool) exists with that name.""" + ts = ToolSet() + schema = ToolSchema(name="plain", description="Plain schema") + ts.add_tool(schema) + assert ts.get_tool("plain") is None + + def test_get_tool_returns_functiontool(self): + """get_tool returns the FunctionTool when present.""" + ts = ToolSet() + tool = FunctionTool(name="found", description="Found tool") + ts.add_tool(tool) + result = ts.get_tool("found") + assert result is tool + + def test_normalize_sorts_by_name(self): + """normalize() sorts tools alphabetically by name.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="z", description="Z")) + ts.add_tool(FunctionTool(name="a", description="A")) + ts.add_tool(FunctionTool(name="m", description="M")) + ts.normalize() + assert [t.name for t in ts.tools] == ["a", "m", "z"] + + def test_names(self): + """names() returns list of tool names.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="a", description="A")) + ts.add_tool(FunctionTool(name="b", description="B")) + assert ts.names() == ["a", "b"] + + def test_func_list(self): + """func_list only includes FunctionTool instances.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="func", description="Func")) + ts.add_tool(ToolSchema(name="plain", description="Plain")) + assert len(ts.func_list) == 1 + assert ts.func_list[0].name == "func" + + def test_merge(self): + """merge() combines tools from another ToolSet.""" + ts1 = ToolSet() + ts1.add_tool(FunctionTool(name="a", description="A")) + ts2 = ToolSet() + ts2.add_tool(FunctionTool(name="b", description="B")) + ts1.merge(ts2) + assert len(ts1) == 2 + + def test_iteration(self): + """ToolSet is iterable and yields ToolSchema items.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="a", description="A")) + ts.add_tool(FunctionTool(name="b", description="B")) + names = [t.name for t in ts] + assert names == ["a", "b"] + + def test_get_light_tool_set(self): + """get_light_tool_set returns tools with empty parameters and no handler.""" + params = { + "type": "object", + "properties": {"q": {"type": "string"}}, + } + ts = ToolSet() + ts.add_tool(FunctionTool( + name="search", description="Search", parameters=params, + handler=AsyncMock(), + )) + light = ts.get_light_tool_set() + assert len(light) == 1 + light_tool = light.get_tool("search") + assert light_tool is not None + assert light_tool.description == "Search" + assert light_tool.parameters == {"type": "object", "properties": {}} + assert light_tool.handler is None + + def test_get_light_tool_skips_inactive(self): + """get_light_tool_set skips inactive tools.""" + ts = ToolSet() + ts.add_tool(FunctionTool(name="active", description="Active", active=True)) + ts.add_tool(FunctionTool(name="inactive", description="Inactive", active=False)) + light = ts.get_light_tool_set() + assert light.names() == ["active"] + + def test_get_param_only_tool_set(self): + """get_param_only_tool_set returns tools with empty description.""" + params = { + "type": "object", + "properties": {"x": {"type": "integer"}}, + } + ts = ToolSet() + ts.add_tool(FunctionTool(name="calc", description="Calc", parameters=params)) + param_only = ts.get_param_only_tool_set() + assert len(param_only) == 1 + tool = param_only.get_tool("calc") + assert tool is not None + assert tool.description == "" + assert tool.parameters == params + + def test_anthropic_schema(self): + """anthropic_schema returns Anthropic-compatible format.""" + ts = ToolSet() + params = { + "type": "object", + "properties": {"x": {"type": "integer"}}, + "required": ["x"], + } + ts.add_tool(FunctionTool(name="add", description="Add", parameters=params)) + schema = ts.anthropic_schema() + assert len(schema) == 1 + assert schema[0]["name"] == "add" + assert schema[0]["description"] == "Add" + assert schema[0]["input_schema"]["properties"] == {"x": {"type": "integer"}} + assert schema[0]["input_schema"]["required"] == ["x"] + + def test_google_schema(self): + """google_schema returns Google GenAI-compatible format.""" + ts = ToolSet() + params = { + "type": "object", + "properties": {"val": {"type": "number", "description": "A value"}}, + } + ts.add_tool(FunctionTool(name="sqrt", description="Square root", parameters=params)) + schema = ts.google_schema() + assert "function_declarations" in schema + assert schema["function_declarations"][0]["name"] == "sqrt" diff --git a/tests/unit/test_tool_executor.py b/tests/unit/test_tool_executor.py new file mode 100644 index 000000000..eda20a5a3 --- /dev/null +++ b/tests/unit/test_tool_executor.py @@ -0,0 +1,136 @@ +"""Unit tests for astrbot.core.agent.tool_executor: BaseFunctionToolExecutor abstract signature.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.tool import FunctionTool +from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor + + +class TestBaseFunctionToolExecutor: + """BaseFunctionToolExecutor cannot be instantiated; metachecks on the ABC.""" + + def test_cannot_instantiate_abstract(self): + """BaseFunctionToolExecutor cannot be instantiated directly (abstract).""" + with pytest.raises(TypeError, match="abstract"): + BaseFunctionToolExecutor() + + def test_execute_is_abstract_classmethod(self): + """execute is both a classmethod and abstractmethod.""" + assert hasattr(BaseFunctionToolExecutor.execute, "__isabstractmethod__") + assert BaseFunctionToolExecutor.execute.__isabstractmethod__ is True + + def test_execute_is_classmethod_descriptor(self): + """execute is a classmethod (has __func__).""" + assert hasattr(BaseFunctionToolExecutor.execute, "__func__") + + def test_concrete_subclass_must_implement_execute(self): + """A subclass without execute is still abstract.""" + class Missing(BaseFunctionToolExecutor): + pass + with pytest.raises(TypeError, match="abstract"): + Missing() + + def test_concrete_subclass_with_execute_can_instantiate(self): + """A subclass that implements execute can be instantiated.""" + class Concrete(BaseFunctionToolExecutor): + @classmethod + async def execute(cls, tool, run_context, **tool_args): + yield "result" + + instance = Concrete() + assert isinstance(instance, BaseFunctionToolExecutor) + + def test_execute_signature_matches(self): + """execute has the expected parameter names.""" + import inspect + sig = inspect.signature( + BaseFunctionToolExecutor.__dict__["execute"].__func__ + ) + param_names = list(sig.parameters.keys()) + assert "tool" in param_names + assert "run_context" in param_names + + def test_execute_tool_parameter_type_hint(self): + """The tool parameter is annotated as FunctionTool.""" + import inspect + sig = inspect.signature( + BaseFunctionToolExecutor.__dict__["execute"].__func__ + ) + tool_param = sig.parameters["tool"] + assert tool_param.annotation is FunctionTool + + def test_execute_run_context_type_hint(self): + """The run_context parameter is annotated as ContextWrapper.""" + import inspect + sig = inspect.signature( + BaseFunctionToolExecutor.__dict__["execute"].__func__ + ) + ctx_param = sig.parameters["run_context"] + origin = getattr(ctx_param.annotation, "__origin__", None) + assert origin is ContextWrapper + + def test_execute_returns_async_generator(self): + """execute return annotation is AsyncGenerator.""" + import inspect + sig = inspect.signature( + BaseFunctionToolExecutor.__dict__["execute"].__func__ + ) + return_annotation = sig.return_annotation + assert "AsyncGenerator" in str(return_annotation) + + def test_concrete_subclass_execute_yields(self): + """A concrete subclass can actually yield values.""" + class Tester(BaseFunctionToolExecutor): + @classmethod + async def execute(cls, tool, run_context, **tool_args): + yield "step1" + yield "step2" + + import asyncio + tool = FunctionTool(name="test", description="test") + ctx = ContextWrapper(context="test") + gen = Tester.execute(tool, ctx) + results = asyncio.run(async_collect(gen)) + assert results == ["step1", "step2"] + + def test_concrete_subclass_generic_parameter(self): + """Subclass can specialize the generic type parameter.""" + class TypedExecutor(BaseFunctionToolExecutor[str]): + @classmethod + async def execute(cls, tool, run_context, **tool_args): + yield run_context.context + + import asyncio + tool = FunctionTool(name="t", description="t") + ctx = ContextWrapper(context="hello_generic") + gen = TypedExecutor.execute(tool, ctx) + results = asyncio.run(async_collect(gen)) + assert results == ["hello_generic"] + + def test_subclass_with_kwargs_passthrough(self): + """execute passes **tool_args through.""" + class KwargsExecutor(BaseFunctionToolExecutor): + @classmethod + async def execute(cls, tool, run_context, **tool_args): + yield tool_args + + import asyncio + tool = FunctionTool(name="t", description="t") + ctx = ContextWrapper(context="ctx") + gen = KwargsExecutor.execute(tool, ctx, x=1, y="two") + results = asyncio.run(async_collect(gen)) + assert results[0] == {"x": 1, "y": "two"} + + +async def async_collect(async_gen): + """Collect all items from an async generator into a list.""" + results = [] + async for item in async_gen: + results.append(item) + return results