ForcePilot/backend/test/unit/channel/lifecycle/test_manager.py
Kris bab30f2715
Some checks failed
Deploy VitePress site to Pages / build (push) Has been cancelled
Ruff Format Check / Ruff Format & Lint (push) Has been cancelled
Deploy VitePress site to Pages / Deploy (push) Has been cancelled
feat:0715
2026-07-15 12:30:58 +08:00

673 lines
25 KiB
Python

from __future__ import annotations
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from yuxi.channel.exceptions import ChannelTransportReconnectRequested
from yuxi.channel.lifecycle.manager import ChannelInstance, ChannelLifecycleManager, _build_webhook_callback_url
from yuxi.channel.plugins.protocol import ChannelMeta, TransportType
from yuxi.channel.transport.protocol import TransportState
class _DummyTask:
"""可 await 且可被 cancel 的任务占位符。"""
def __init__(self):
self.cancel = MagicMock()
def __await__(self):
return iter([])
class TestBuildWebhookCallbackUrl:
def test_uses_default_base_url(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("CHANNEL_WEBHOOK_CALLBACK_URL", raising=False)
assert _build_webhook_callback_url("wechat") == "http://localhost:5050/api/channels/wechat/webhook"
def test_uses_env_var(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("CHANNEL_WEBHOOK_CALLBACK_URL", "https://example.com/")
assert _build_webhook_callback_url("wechat") == "https://example.com/api/channels/wechat/webhook"
def test_strips_trailing_slash(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("CHANNEL_WEBHOOK_CALLBACK_URL", "https://example.com")
assert _build_webhook_callback_url("wechat") == "https://example.com/api/channels/wechat/webhook"
class TestChannelInstance:
def test_initial_state(self):
plugin = MagicMock()
instance = ChannelInstance("wechat", "acc-1", {"token": "x"}, plugin)
assert instance.channel_type == "wechat"
assert instance.account_id == "acc-1"
assert instance.config == {"token": "x"}
assert instance.plugin is plugin
assert instance.transport is None
assert instance.no_transport is False
assert instance.last_error is None
assert instance.reconnect_attempts == 0
assert instance.is_closed() is False
async def test_start_creates_transport(self):
plugin = MagicMock()
plugin.supports_qr_login.return_value = False
transport = MagicMock()
transport.state = TransportState.CONNECTED
transport.start = AsyncMock()
transport.on_message = MagicMock()
plugin.create_transport = AsyncMock(return_value=transport)
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
await instance.start()
assert instance.transport is transport
assert instance.no_transport is False
transport.start.assert_awaited_once()
async def test_start_with_no_transport(self):
plugin = MagicMock()
plugin.create_transport = AsyncMock(return_value=None)
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
await instance.start()
assert instance.transport is None
assert instance.no_transport is True
async def test_start_stops_existing_transport(self):
plugin = MagicMock()
plugin.supports_qr_login.return_value = False
old_transport = MagicMock()
old_transport.state = TransportState.CONNECTED
old_transport.stop = AsyncMock()
new_transport = MagicMock()
new_transport.state = TransportState.CONNECTED
new_transport.start = AsyncMock()
new_transport.on_message = MagicMock()
plugin.create_transport = AsyncMock(return_value=new_transport)
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance._transport = old_transport
await instance.start()
old_transport.stop.assert_awaited_once()
assert instance.transport is new_transport
async def test_stop(self):
plugin = MagicMock()
transport = MagicMock()
transport.stop = AsyncMock()
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance._transport = transport
await instance.stop(timeout=5.0)
transport.stop.assert_awaited_once()
assert instance.transport is None
async def test_stop_without_transport_is_safe(self):
instance = ChannelInstance("wechat", "acc-1", {}, MagicMock())
await instance.stop()
assert instance.transport is None
async def test_stop_times_out(self):
plugin = MagicMock()
transport = MagicMock()
async def slow_stop():
await asyncio.sleep(10)
transport.stop = slow_stop
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance._transport = transport
await instance.stop(timeout=0.01)
assert instance.transport is None
def test_set_message_handler_attaches_to_existing_transport(self):
plugin = MagicMock()
transport = MagicMock()
transport.on_message = MagicMock()
handler = AsyncMock()
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance._transport = transport
instance.set_message_handler(handler)
assert instance._message_handler is handler
transport.on_message.assert_called_once()
def test_set_message_handler_without_transport(self):
plugin = MagicMock()
handler = AsyncMock()
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance.set_message_handler(handler)
assert instance._message_handler is handler
async def test_register_transport_message_handler(self):
plugin = MagicMock()
transport = MagicMock()
transport.on_message = MagicMock()
handler = AsyncMock()
instance = ChannelInstance("wechat", "acc-1", {}, plugin, message_handler=handler)
instance._transport = transport
instance._register_transport_message_handler()
transport.on_message.assert_called_once()
wrapped = transport.on_message.call_args[0][0]
raw = b"hello"
await wrapped(raw)
# Wait briefly for the spawned task
await asyncio.sleep(0)
await asyncio.sleep(0)
handler.assert_awaited_once_with(raw, "wechat", "acc-1")
async def test_register_transport_message_handler_reconnect_requested(self):
plugin = MagicMock()
transport = MagicMock()
transport.on_message = MagicMock()
transport.stop = AsyncMock()
handler = AsyncMock(side_effect=ChannelTransportReconnectRequested("reconnect"))
instance = ChannelInstance("wechat", "acc-1", {}, plugin, message_handler=handler)
instance._transport = transport
instance._register_transport_message_handler()
wrapped = transport.on_message.call_args[0][0]
await wrapped(b"hello")
await asyncio.sleep(0)
await asyncio.sleep(0)
handler.assert_awaited_once()
transport.stop.assert_awaited_once()
assert instance._reconnect_requested is True
async def test_register_transport_message_handler_error(self):
plugin = MagicMock()
transport = MagicMock()
transport.on_message = MagicMock()
handler = AsyncMock(side_effect=RuntimeError("boom"))
instance = ChannelInstance("wechat", "acc-1", {}, plugin, message_handler=handler)
instance._transport = transport
instance._register_transport_message_handler()
wrapped = transport.on_message.call_args[0][0]
await wrapped(b"hello")
await asyncio.sleep(0)
await asyncio.sleep(0)
handler.assert_awaited_once()
def test_record_error(self):
instance = ChannelInstance("wechat", "acc-1", {}, MagicMock())
instance.record_error(RuntimeError("boom"))
assert instance.last_error == "boom"
assert instance.reconnect_attempts == 1
async def test_wait_for_exit_with_transport(self):
plugin = MagicMock()
transport = MagicMock()
transport.state = TransportState.CONNECTED
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance._transport = transport
async def waiter():
await asyncio.sleep(0.05)
transport.state = TransportState.DISCONNECTED
task = asyncio.create_task(waiter())
await instance.wait_for_exit()
await task
async def test_wait_for_exit_without_transport(self):
instance = ChannelInstance("wechat", "acc-1", {}, MagicMock())
async def closer():
await asyncio.sleep(0.05)
instance.close()
task = asyncio.create_task(closer())
await instance.wait_for_exit()
await task
def test_close(self):
instance = ChannelInstance("wechat", "acc-1", {}, MagicMock())
instance.close()
assert instance.is_closed() is True
class TestChannelLifecycleManager:
def _make_manager(self, registry=None, config_manager=None, handler=None):
registry = registry or MagicMock()
config_manager = config_manager or MagicMock()
return ChannelLifecycleManager(registry, config_manager, handler)
def test_initial_state(self):
registry = MagicMock()
config_manager = MagicMock()
manager = ChannelLifecycleManager(registry, config_manager)
assert manager.registry is registry
assert manager.config_manager is config_manager
assert manager._instances == {}
assert manager._tasks == {}
def test_instance_key(self):
manager = self._make_manager()
assert manager._instance_key("wechat", "acc-1") == "wechat:acc-1"
async def test_set_message_handler_propagates_to_instances(self):
manager = self._make_manager()
instance = MagicMock()
instance.set_message_handler = MagicMock()
manager._instances["wechat:acc-1"] = instance
handler = AsyncMock()
manager.set_message_handler(handler)
assert manager._message_handler is handler
instance.set_message_handler.assert_called_once_with(handler)
async def test_start_all_starts_enabled_accounts(self):
config_manager = MagicMock()
config_manager.list_enabled_accounts = AsyncMock(
return_value=[
{"channel_type": "wechat", "account_id": "acc-1"},
{"channel_type": "slack", "account_id": "acc-2"},
]
)
manager = self._make_manager(config_manager=config_manager)
manager.start_channel = AsyncMock()
await manager.start_all()
manager.start_channel.assert_any_await("wechat", "acc-1")
manager.start_channel.assert_any_await("slack", "acc-2")
assert manager.start_channel.await_count == 2
async def test_start_with_semaphore_uses_semaphore(self):
manager = self._make_manager()
manager.start_channel = AsyncMock()
config = {"channel_type": "wechat", "account_id": "acc-1"}
await manager._start_with_semaphore(config)
manager.start_channel.assert_awaited_once_with("wechat", "acc-1")
async def test_start_channel_skips_existing(self):
manager = self._make_manager()
manager._instances["wechat:acc-1"] = MagicMock()
await manager.start_channel("wechat", "acc-1")
manager.registry.get_plugin.assert_not_called()
async def test_start_channel_warns_when_plugin_missing(self):
manager = self._make_manager()
manager.registry.get_plugin = MagicMock(return_value=None)
await manager.start_channel("unknown", "acc-1")
manager.registry.get_plugin.assert_called_once_with("unknown")
async def test_start_channel_handles_config_load_error(self):
config_manager = MagicMock()
config_manager.get_config = AsyncMock(side_effect=RuntimeError("db down"))
registry = MagicMock()
registry.get_plugin = MagicMock(return_value=MagicMock())
manager = self._make_manager(registry=registry, config_manager=config_manager)
await manager.start_channel("wechat", "acc-1")
assert "wechat:acc-1" not in manager._instances
async def test_start_channel_creates_instance(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.POLLING
)
plugin.run_startup_maintenance = AsyncMock()
plugin.on_channel_enabled = AsyncMock()
config_manager = MagicMock()
config_manager.get_config = AsyncMock(return_value={"token": "x"})
registry = MagicMock()
registry.get_plugin = MagicMock(return_value=plugin)
manager = self._make_manager(registry=registry, config_manager=config_manager)
with patch("yuxi.channel.lifecycle.manager.asyncio.create_task") as create_task:
dummy_task = MagicMock()
create_task.return_value = dummy_task
await manager.start_channel("wechat", "acc-1")
assert "wechat:acc-1" in manager._instances
instance = manager._instances["wechat:acc-1"]
assert instance.channel_type == "wechat"
assert instance.account_id == "acc-1"
plugin.run_startup_maintenance.assert_awaited_once_with({"token": "x"}, "acc-1")
plugin.on_channel_enabled.assert_awaited_once_with({"token": "x"}, "acc-1")
async def test_start_channel_webhook_setup(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBHOOK
)
plugin.run_startup_maintenance = AsyncMock()
plugin.setup_webhook = AsyncMock(return_value=True)
plugin.on_channel_enabled = AsyncMock()
config_manager = MagicMock()
config_manager.get_config = AsyncMock(return_value={"token": "x"})
registry = MagicMock()
registry.get_plugin = MagicMock(return_value=plugin)
manager = self._make_manager(registry=registry, config_manager=config_manager)
with patch("yuxi.channel.lifecycle.manager.asyncio.create_task"):
await manager.start_channel("wechat", "acc-1")
plugin.setup_webhook.assert_awaited_once()
plugin.on_channel_enabled.assert_awaited_once()
async def test_start_channel_webhook_setup_failure_is_logged(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBHOOK
)
plugin.run_startup_maintenance = AsyncMock()
plugin.setup_webhook = AsyncMock(side_effect=RuntimeError("setup failed"))
plugin.on_channel_enabled = AsyncMock()
config_manager = MagicMock()
config_manager.get_config = AsyncMock(return_value={"token": "x"})
registry = MagicMock()
registry.get_plugin = MagicMock(return_value=plugin)
manager = self._make_manager(registry=registry, config_manager=config_manager)
with patch("yuxi.channel.lifecycle.manager.asyncio.create_task"):
await manager.start_channel("wechat", "acc-1")
assert "wechat:acc-1" in manager._instances
plugin.on_channel_enabled.assert_awaited_once()
async def test_stop_channel(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.POLLING
)
plugin.on_channel_disabled = AsyncMock()
instance = MagicMock()
instance.channel_type = "wechat"
instance.account_id = "acc-1"
instance.config = {"token": "x"}
instance.plugin = plugin
instance.close = MagicMock()
instance.stop = AsyncMock()
manager = self._make_manager()
manager._instances["wechat:acc-1"] = instance
task = _DummyTask()
manager._tasks["wechat:acc-1"] = task
await manager.stop_channel("wechat", "acc-1")
instance.close.assert_called_once()
instance.stop.assert_awaited_once_with(timeout=5.0)
task.cancel.assert_called_once()
assert "wechat:acc-1" not in manager._instances
plugin.on_channel_disabled.assert_awaited_once_with({"token": "x"}, "acc-1")
async def test_stop_channel_webhook_delete(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBHOOK
)
plugin.delete_webhook = AsyncMock(return_value=True)
plugin.on_channel_disabled = AsyncMock()
instance = MagicMock()
instance.channel_type = "wechat"
instance.account_id = "acc-1"
instance.config = {"token": "x"}
instance.plugin = plugin
instance.close = MagicMock()
instance.stop = AsyncMock()
manager = self._make_manager()
manager._instances["wechat:acc-1"] = instance
task = _DummyTask()
manager._tasks["wechat:acc-1"] = task
await manager.stop_channel("wechat", "acc-1")
plugin.delete_webhook.assert_awaited_once_with({"token": "x"})
plugin.on_channel_disabled.assert_awaited_once()
async def test_stop_channel_missing_instance(self):
manager = self._make_manager()
await manager.stop_channel("wechat", "acc-1")
def test_get_transport_returns_instance_transport(self):
manager = self._make_manager()
instance = MagicMock()
instance.transport = MagicMock()
manager._instances["wechat:acc-1"] = instance
transport = manager.get_transport("wechat", "acc-1")
assert transport is instance.transport
def test_get_transport_returns_none_when_instance_missing(self):
manager = self._make_manager()
assert manager.get_transport("wechat", "acc-1") is None
async def test_stop_all(self):
manager = self._make_manager()
manager._instances["wechat:acc-1"] = MagicMock(channel_type="wechat", account_id="acc-1")
manager._instances["slack:acc-2"] = MagicMock(channel_type="slack", account_id="acc-2")
manager.stop_channel = AsyncMock()
await manager.stop_all()
assert manager.stop_channel.await_count == 2
async def test_run_loop_starts_transport_and_waits(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBSOCKET
)
transport = MagicMock()
transport.state = TransportState.CONNECTED
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance.start = AsyncMock()
instance._no_transport = False
instance._transport = transport
instance.wait_for_exit = AsyncMock()
instance.stop = AsyncMock()
manager = self._make_manager()
with patch("yuxi.channel.lifecycle.manager.channel_active_connections") as mock_gauge:
with patch("yuxi.channel.lifecycle.manager.channel_reconnect_total") as mock_counter:
run_task = asyncio.create_task(manager._run_loop(instance))
await asyncio.sleep(0)
instance.close()
try:
await asyncio.wait_for(run_task, timeout=0.2)
except TimeoutError:
run_task.cancel()
try:
await run_task
except asyncio.CancelledError:
pass
instance.start.assert_awaited()
instance.wait_for_exit.assert_awaited()
async def test_run_loop_no_transport_breaks(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBHOOK
)
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance.start = AsyncMock()
instance._no_transport = True
instance._transport = None
instance.wait_for_exit = AsyncMock()
instance.stop = AsyncMock()
manager = self._make_manager()
with patch("yuxi.channel.lifecycle.manager.channel_active_connections") as mock_gauge:
await manager._run_loop(instance)
instance.start.assert_awaited_once()
instance.wait_for_exit.assert_not_awaited()
async def test_run_loop_reconnects_after_disconnect(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBSOCKET
)
transport = MagicMock()
transport.state = TransportState.CONNECTED
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
call_count = {"start": 0}
async def start_side_effect():
call_count["start"] += 1
instance._no_transport = False
instance._transport = transport
instance.start = AsyncMock(side_effect=start_side_effect)
instance._no_transport = False
instance._transport = transport
wait_count = {"n": 0}
async def wait_side_effect():
wait_count["n"] += 1
if wait_count["n"] == 1:
transport.state = TransportState.DISCONNECTED
return
instance.close()
instance.wait_for_exit = AsyncMock(side_effect=wait_side_effect)
instance.stop = AsyncMock()
manager = self._make_manager()
with patch("yuxi.channel.lifecycle.manager.channel_active_connections") as mock_gauge:
with patch("yuxi.channel.lifecycle.manager.channel_reconnect_total") as mock_counter:
with patch("yuxi.channel.lifecycle.manager.asyncio.sleep", new=AsyncMock()):
await manager._run_loop(instance)
assert call_count["start"] == 2
async def test_run_loop_respects_cancel(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBSOCKET
)
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance.start = AsyncMock(side_effect=asyncio.CancelledError)
instance.stop = AsyncMock()
manager = self._make_manager()
await manager._run_loop(instance)
instance.start.assert_awaited_once()
async def test_run_loop_records_error_and_backoff(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBSOCKET
)
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance.start = AsyncMock(side_effect=ConnectionError("boom"))
instance.stop = AsyncMock()
instance.record_error = MagicMock(wraps=instance.record_error)
instance.is_closed = MagicMock(side_effect=[False, True])
manager = self._make_manager()
with patch("yuxi.channel.lifecycle.manager.channel_reconnect_total") as mock_counter:
with patch("yuxi.channel.lifecycle.manager.asyncio.sleep", new=AsyncMock()):
await manager._run_loop(instance)
instance.record_error.assert_called_once()
assert instance.reconnect_attempts == 1
async def test_run_loop_reconnect_requested_skips_backoff(self):
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.WEBSOCKET
)
transport = MagicMock()
transport.state = TransportState.CONNECTED
instance = ChannelInstance("wechat", "acc-1", {}, plugin)
instance.start = AsyncMock()
instance._no_transport = False
instance._transport = transport
instance.stop = AsyncMock()
instance.record_error = MagicMock(wraps=instance.record_error)
wait_count = {"n": 0}
async def wait_side_effect():
wait_count["n"] += 1
if wait_count["n"] == 1:
instance._reconnect_requested = True
return
instance.close()
instance.wait_for_exit = AsyncMock(side_effect=wait_side_effect)
manager = self._make_manager()
with patch("yuxi.channel.lifecycle.manager.channel_active_connections"):
with patch("yuxi.channel.lifecycle.manager.channel_reconnect_total") as mock_counter:
await manager._run_loop(instance)
instance.record_error.assert_not_called()
mock_counter.inc.assert_not_called()
assert instance.context.reconnect_attempts == 0
async def test_per_key_lock_allows_concurrent_start(self):
"""不同 key 的 start_channel 可以并发执行。"""
plugin = MagicMock()
plugin.get_meta.return_value = ChannelMeta(
channel_type="wechat", display_name="WeChat", transport_type=TransportType.POLLING
)
plugin.run_startup_maintenance = AsyncMock()
plugin.on_channel_enabled = AsyncMock()
config_manager = MagicMock()
config_manager.get_config = AsyncMock(return_value={"token": "x"})
registry = MagicMock()
registry.get_plugin = MagicMock(return_value=plugin)
manager = self._make_manager(registry=registry, config_manager=config_manager)
with patch("yuxi.channel.lifecycle.manager.asyncio.create_task"):
await asyncio.gather(
manager.start_channel("wechat", "acc-1"),
manager.start_channel("slack", "acc-2"),
)
assert "wechat:acc-1" in manager._instances
assert "slack:acc-2" in manager._instances