From 9d03844371a790a2123ae8487bcfe69383219242 Mon Sep 17 00:00:00 2001 From: Jack Kingsman Date: Fri, 6 Mar 2026 14:09:30 -0800 Subject: [PATCH] Improve module lifecycling --- app/fanout/bot.py | 17 ++- app/fanout/manager.py | 1 + tests/test_fanout_integration.py | 214 +++++++++++++++++++++++++++++++ 3 files changed, 230 insertions(+), 2 deletions(-) diff --git a/app/fanout/bot.py b/app/fanout/bot.py index 36c49e9..5b6d8e0 100644 --- a/app/fanout/bot.py +++ b/app/fanout/bot.py @@ -21,10 +21,23 @@ class BotModule(FanoutModule): def __init__(self, config_id: str, config: dict, *, name: str = "Bot") -> None: super().__init__(config_id, config, name=name) + self._tasks: set[asyncio.Task] = set() + self._active = True + + async def stop(self) -> None: + self._active = False + for task in self._tasks: + task.cancel() + # Wait briefly for tasks to acknowledge cancellation + if self._tasks: + await asyncio.gather(*self._tasks, return_exceptions=True) + self._tasks.clear() async def on_message(self, data: dict) -> None: """Kick off bot execution in a background task so we don't block dispatch.""" - asyncio.create_task(self._run_for_message(data)) + task = asyncio.create_task(self._run_for_message(data)) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) async def _run_for_message(self, data: dict) -> None: from app.fanout.bot_exec import ( @@ -118,7 +131,7 @@ class BotModule(FanoutModule): logger.warning("Bot '%s' execution error: %s", self.name, e) return - if response: + if response and self._active: await process_bot_response(response, is_dm, sender_key or "", channel_key) @property diff --git a/app/fanout/manager.py b/app/fanout/manager.py index b58041b..665179a 100644 --- a/app/fanout/manager.py +++ b/app/fanout/manager.py @@ -194,6 +194,7 @@ class FanoutManager: await module.start() except Exception: logger.exception("Failed to restart timed-out fanout module %s", config_id) + self._modules.pop(config_id, None) async def broadcast_message(self, data: dict) -> None: """Dispatch a decoded message to modules whose scope matches.""" diff --git a/tests/test_fanout_integration.py b/tests/test_fanout_integration.py index 42c7e68..22df257 100644 --- a/tests/test_fanout_integration.py +++ b/tests/test_fanout_integration.py @@ -1235,3 +1235,217 @@ class TestFanoutAppriseIntegration: body_text = str(results[0]) assert "Eve" in body_text assert "routed msg" in body_text + + +# --------------------------------------------------------------------------- +# Bot lifecycle tests +# --------------------------------------------------------------------------- + + +class TestBotModuleLifecycle: + """BotModule.stop() must cancel in-flight tasks and prevent response delivery.""" + + @pytest.mark.asyncio + async def test_stop_cancels_pending_tasks(self): + """Stopping a bot module cancels tasks still in the settle delay.""" + from app.fanout.bot import BotModule + + mod = BotModule("bot1", {"code": "def bot(**k): return 'hi'"}, name="Test Bot") + mod._active = True + + # Fire off a message — it will enter the 2s settle sleep + await mod.on_message( + {"type": "PRIV", "conversation_key": "abc123", "text": "hello", "outgoing": False} + ) + assert len(mod._tasks) == 1 + + # Stop immediately — should cancel the pending task + await mod.stop() + + assert mod._active is False + assert len(mod._tasks) == 0 + + @pytest.mark.asyncio + async def test_stop_prevents_response_delivery(self): + """Even if bot code returns a response, stop() prevents it from being sent.""" + from unittest.mock import AsyncMock, patch + + from app.fanout.bot import BotModule + + mod = BotModule("bot1", {"code": "def bot(**k): return 'reply'"}, name="Test Bot") + + mock_process = AsyncMock() + with patch("app.fanout.bot.asyncio.sleep", new_callable=AsyncMock): + # Manually run the handler with _active=True, then set _active=False + # before process_bot_response would be called + original_run = mod._run_for_message + + async def run_then_deactivate(data): + # Deactivate mid-flight by stopping + mod._active = False + await original_run(data) + + with patch.object(mod, "_run_for_message", run_then_deactivate): + await mod.on_message( + { + "type": "PRIV", + "conversation_key": "abc123", + "text": "hello", + "outgoing": False, + } + ) + # Wait for the task to finish + if mod._tasks: + await asyncio.gather(*mod._tasks, return_exceptions=True) + + # process_bot_response should never have been called + mock_process.assert_not_called() + + @pytest.mark.asyncio + async def test_multiple_tasks_all_cancelled(self): + """Multiple in-flight tasks are all cancelled on stop.""" + from app.fanout.bot import BotModule + + mod = BotModule("bot1", {"code": "def bot(**k): return 'hi'"}, name="Test Bot") + mod._active = True + + # Fire off several messages + for i in range(5): + await mod.on_message( + { + "type": "PRIV", + "conversation_key": f"key{i}", + "text": f"msg{i}", + "outgoing": False, + } + ) + assert len(mod._tasks) == 5 + + await mod.stop() + + assert mod._active is False + assert len(mod._tasks) == 0 + + +# --------------------------------------------------------------------------- +# Manager restart failure tests +# --------------------------------------------------------------------------- + + +class TestManagerRestartFailure: + """_restart_module removes dead module from dispatch table on failure.""" + + @pytest.mark.asyncio + async def test_failed_restart_removes_module(self): + """When module.start() fails during restart, the module is removed from _modules.""" + + from app.fanout.base import FanoutModule + + class FailingModule(FanoutModule): + def __init__(self): + super().__init__("fail1", {}, name="Failer") + self.stop_called = False + self.start_calls = 0 + + async def stop(self): + self.stop_called = True + + async def start(self): + self.start_calls += 1 + raise ConnectionError("broker down") + + @property + def status(self): + return "error" + + manager = FanoutManager() + mod = FailingModule() + manager._modules["fail1"] = (mod, {"messages": "all", "raw_packets": "none"}) + + # Restart should catch the error and remove the module + await manager._restart_module("fail1", mod) + + assert mod.stop_called + assert mod.start_calls == 1 + assert "fail1" not in manager._modules + + @pytest.mark.asyncio + async def test_successful_restart_keeps_module(self): + """When restart succeeds, the module stays in _modules.""" + from app.fanout.base import FanoutModule + + class GoodModule(FanoutModule): + def __init__(self): + super().__init__("good1", {}, name="Goodie") + + async def stop(self): + pass + + async def start(self): + pass + + @property + def status(self): + return "connected" + + manager = FanoutManager() + mod = GoodModule() + scope = {"messages": "all", "raw_packets": "none"} + manager._modules["good1"] = (mod, scope) + + await manager._restart_module("good1", mod) + + assert "good1" in manager._modules + + @pytest.mark.asyncio + async def test_dead_module_not_dispatched_after_failed_restart(self): + """After failed restart, the dead module does not receive further dispatches.""" + from app.fanout.base import FanoutModule + + class TrackingModule(FanoutModule): + def __init__(self, config_id): + super().__init__(config_id, {}, name=config_id) + self.messages_received = [] + + async def start(self): + raise RuntimeError("can't start") + + async def stop(self): + pass + + async def on_message(self, data): + self.messages_received.append(data) + + @property + def status(self): + return "error" + + class HealthyModule(FanoutModule): + def __init__(self): + super().__init__("healthy", {}, name="Healthy") + self.messages_received = [] + + async def on_message(self, data): + self.messages_received.append(data) + + @property + def status(self): + return "connected" + + manager = FanoutManager() + dead = TrackingModule("dead1") + healthy = HealthyModule() + + scope = {"messages": "all", "raw_packets": "none"} + manager._modules["dead1"] = (dead, scope) + manager._modules["healthy"] = (healthy, scope) + + # Simulate failed restart of dead module + await manager._restart_module("dead1", dead) + assert "dead1" not in manager._modules + + # Now broadcast — only the healthy module should receive + await manager.broadcast_message({"type": "PRIV", "conversation_key": "k1", "text": "hi"}) + + assert len(healthy.messages_received) == 1 + assert len(dead.messages_received) == 0