Separate staff roles and protect message delivery with operator guides

tests/test_delivery_resilience.py

@@ -0,0 +1,126 @@
+import asyncio
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+from aiogram import Bot
+from aiogram.exceptions import TelegramForbiddenError, TelegramRetryAfter
+from aiogram.methods import CopyMessage
+from aiogram.types import Chat, Message, User as TelegramUser
+
+from src.core.broadcast_services import BroadcastService
+from src.core.database import async_session_maker
+from src.core.services import UserService
+from src.handlers import chat_handlers
+from src.utils import delivery, telegram_messages
+from src.utils.broadcast_jobs import jobs, shutdown_broadcasts, start_broadcast
+from test_dispatcher import TelegramStub, dispatch, event
+
+
+async def test_a_stalled_chat_recipient_does_not_block_other_recipients_or_next_message(monkeypatch):
+    monkeypatch.setattr(telegram_messages, "COPY_TIMEOUT", 0.1)
+    async with async_session_maker() as session:
+        sender = await UserService.get_or_create_user(session, 831)
+        await UserService.get_or_create_user(session, 832)
+        await UserService.get_or_create_user(session, 833)
+
+    class StallingTelegram(TelegramStub):
+        async def make_request(self, bot, method, timeout=None):
+            if getattr(method, "chat_id", None) == 832:
+                await asyncio.Event().wait()
+            return await super().make_request(bot, method, timeout)
+
+    stub = StallingTelegram()
+    bot = Bot("123456:TEST_TOKEN_FOR_ISOLATED_TESTS", session=stub)
+    message = Message(message_id=1, date=datetime.now(timezone.utc), chat=Chat(id=831, type="private"),
+                      from_user=TelegramUser(id=831, first_name="Sender", is_bot=False), text="hello").as_(bot)
+    _, successes, failures = await asyncio.wait_for(
+        chat_handlers.broadcast_message_with_scheduler(message, sender, exclude_user_id=831), 2)
+    assert successes >= 1 and failures == 1
+    assert await chat_handlers._copy_with_sender(message, 833, "Sender") == 100
+
+
+async def test_blocked_recipient_and_extreme_flood_wait_do_not_stop_broadcast(monkeypatch):
+    service = BroadcastService()
+    async with async_session_maker() as session:
+        admin = await UserService.get_or_create_user(session, 900001)
+        users = [await UserService.get_or_create_user(session, number) for number in (834, 835, 836)]
+
+    async def copy_to(target):
+        method = CopyMessage(chat_id=target, from_chat_id=1, message_id=1)
+        if target == 834:
+            raise TelegramForbiddenError(method=method, message="bot was blocked by the user")
+        if target == 835:
+            raise TelegramRetryAfter(method=method, message="slow down", retry_after=3600)
+        return SimpleNamespace(message_id=1)
+
+    message = SimpleNamespace(copy_to=copy_to, text="test", caption=None, content_type="text",
+                              photo=None, video=None, document=None, animation=None, voice=None, audio=None)
+    result = await asyncio.wait_for(service.broadcast_to_users(None, message, admin.id, users), 3)
+    assert result == {"total": 3, "success": 1, "failed": 2, "blocked": 1}
+
+
+async def test_timeout_releases_delivery_capacity(monkeypatch):
+    monkeypatch.setattr(delivery, "DELIVERY_TIMEOUT", 0.05)
+
+    @delivery.background_delivery
+    async def stall():
+        await asyncio.Event().wait()
+
+    result = await asyncio.gather(*(stall() for _ in range(8)), return_exceptions=True)
+    assert all(isinstance(error, TimeoutError) for error in result)
+    sender = SimpleNamespace(send_message=AsyncMock(return_value=True))
+    assert await delivery.send_background(sender, chat_id=1, text="next")
+
+
+async def test_background_job_failure_releases_slot_and_shutdown_cancels_pending_work():
+    failure_report = AsyncMock()
+
+    async def broken():
+        raise RuntimeError("synthetic failure")
+
+    assert start_broadcast(837, broken, failure_report)
+    assert not start_broadcast(837, broken, failure_report)
+    await asyncio.gather(*list(jobs().values()))
+    failure_report.assert_awaited_once()
+    assert not jobs()
+    assert start_broadcast(837, asyncio.Event().wait, failure_report)
+    await shutdown_broadcasts()
+    assert not jobs()
+
+
+async def test_missing_broadcast_report_does_not_fail_completed_delivery():
+    from src.handlers.admin_panel import _edit_broadcast_status
+    report = SimpleNamespace(edit_text=AsyncMock(side_effect=TelegramForbiddenError(
+        method=CopyMessage(chat_id=1, from_chat_id=1, message_id=1), message="bot was blocked")))
+    await _edit_broadcast_status(report, "Delivery completed")
+    report.edit_text.assert_awaited_once()
+
+
+async def test_admin_can_use_start_while_broadcast_runs_and_failure_clears_fsm(monkeypatch):
+    import main
+    from src.handlers import admin_panel
+    started = asyncio.Event()
+    finish = asyncio.Event()
+
+    async def sending(*args, **kwargs):
+        started.set()
+        await finish.wait()
+        raise RuntimeError("synthetic Telegram failure")
+
+    monkeypatch.setattr(admin_panel, "_broadcast_direct", sending)
+    await dispatch(event(900001, text="/cancel"))
+    state = main.dp.fsm.get_context(bot=Bot("123456:TEST_TOKEN_FOR_ISOLATED_TESTS"), chat_id=900001, user_id=900001)
+    await state.set_state(admin_panel.AdminStates.broadcast_message)
+    await state.set_data({"broadcast_type": "direct"})
+    try:
+        await dispatch(event(900001, text="broadcast content"))
+        await asyncio.wait_for(started.wait(), 1)
+        assert await state.get_state() is None
+        calls = await asyncio.wait_for(dispatch(event(900001, text="/start")), 2)
+        assert any((getattr(call, "text", "") or "").startswith("👋") for call in calls)
+        finish.set()
+        await asyncio.gather(*list(jobs().values()))
+        assert await state.get_state() is None
+    finally:
+        await shutdown_broadcasts()