tgit/tests/fakes.py

"""Offline Telegram: a fake aiogram session that records outgoing API calls and fabricates replies."""

from collections.abc import AsyncGenerator, Callable
from datetime import UTC, datetime
from itertools import count
from typing import Any

from aiogram import Bot, Dispatcher
from aiogram.client.session.base import BaseSession
from aiogram.exceptions import TelegramForbiddenError
from aiogram.methods import EditMessageText, GetFile, GetMe, SendMessage, SendPhoto, TelegramMethod
from aiogram.methods.base import TelegramType
from aiogram.types import Chat, File, Message, Update, User

BOT_TOKEN = "42:TEST-TOKEN"


class FakeSession(BaseSession):
    def __init__(self) -> None:
        super().__init__()
        self.requests: list[TelegramMethod[Any]] = []
        self._ids = count(1000)
        self.files: dict[str, bytes] = {}  # file_id -> content, for documents users "upload"
        self.blocked: set[int] = set()  # chats that blocked the bot: sends there raise 403

    async def make_request(
        self,
        bot: Bot,
        method: TelegramMethod[TelegramType],
        timeout: int | None = None,  # noqa: ASYNC109 - signature fixed by aiogram BaseSession
    ) -> TelegramType:
        self.requests.append(method)
        if isinstance(method, SendMessage) and method.chat_id in self.blocked:
            raise TelegramForbiddenError(method=method, message="Forbidden: bot was blocked by the user")
        if isinstance(method, GetFile):
            return File(  # type: ignore[return-value]
                file_id=method.file_id, file_unique_id=method.file_id, file_path=f"documents/{method.file_id}"
            )
        if isinstance(method, GetMe):
            return User(id=42, is_bot=True, first_name="tgit", username="tgit_test_bot")  # type: ignore[return-value]
        if isinstance(method, SendPhoto):
            chat = method.chat_id if isinstance(method.chat_id, int) else 0
            return Message(  # type: ignore[return-value]
                message_id=next(self._ids),
                date=datetime.now(UTC),
                chat=Chat(id=chat, type="private"),
                caption=method.caption,
            )
        if isinstance(method, SendMessage | EditMessageText):
            chat_id = method.chat_id if isinstance(method.chat_id, int) else 0
            return Message(  # type: ignore[return-value]
                message_id=next(self._ids),
                date=datetime.now(UTC),
                chat=Chat(id=chat_id, type="private"),
                text=method.text,
            )
        return True  # type: ignore[return-value]

    async def stream_content(
        self,
        url: str,
        headers: dict[str, Any] | None = None,
        timeout: int = 30,  # noqa: ASYNC109 - signature fixed by aiogram BaseSession
        chunk_size: int = 65536,
        raise_for_status: bool = True,
    ) -> AsyncGenerator[bytes, None]:
        file_id = url.rsplit("/", 1)[-1]
        yield self.files[file_id]

    async def close(self) -> None:
        pass

    def sent(self, method_type: type[TelegramMethod[Any]]) -> list[Any]:
        return [r for r in self.requests if isinstance(r, method_type)]


class TgClient:
    """Feeds updates through a real Dispatcher (filters, middlewares, FSM) on top of FakeSession."""

    def __init__(self, dp: Dispatcher | Callable[[Bot], Dispatcher]) -> None:
        """dp: a dispatcher, or a factory for one that needs the bot (alert broadcasts send through it)."""
        self.session = FakeSession()
        self.bot = Bot(BOT_TOKEN, session=self.session)
        self.dp = dp if isinstance(dp, Dispatcher) else dp(self.bot)
        self._update_ids = count(1)
        self._message_ids = count(1)

    @staticmethod
    def user(tg_id: int) -> dict[str, Any]:
        return {"id": tg_id, "is_bot": False, "first_name": f"u{tg_id}", "username": f"u{tg_id}"}

    async def send(self, tg_id: int, text: str) -> None:
        message = {
            "message_id": next(self._message_ids),
            "date": int(datetime.now(UTC).timestamp()),
            "chat": {"id": tg_id, "type": "private"},
            "from": self.user(tg_id),
            "text": text,
        }
        if text.startswith("/"):
            cmd = text.split()[0]
            message["entities"] = [{"type": "bot_command", "offset": 0, "length": len(cmd)}]
        await self._feed({"message": message})

    async def send_document(self, tg_id: int, name: str, content: bytes) -> None:
        file_id = f"doc{next(self._message_ids)}"
        self.session.files[file_id] = content
        message = {
            "message_id": next(self._message_ids),
            "date": int(datetime.now(UTC).timestamp()),
            "chat": {"id": tg_id, "type": "private"},
            "from": self.user(tg_id),
            "document": {"file_id": file_id, "file_unique_id": file_id, "file_name": name, "file_size": len(content)},
        }
        await self._feed({"message": message})

    async def press_markup(self, tg_id: int, data: str, markup: Any) -> None:
        """Press a button on a message that carries this inline keyboard (some handlers look at it)."""
        await self.press(tg_id, data, reply_markup=markup.model_dump(exclude_none=True) if markup else None)

    async def press(
        self, tg_id: int, data: str, message_text: str = "section", reply_markup: dict[str, Any] | None = None
    ) -> None:
        callback = {
            "id": str(next(self._update_ids)),
            "from": self.user(tg_id),
            "chat_instance": "ci",
            "data": data,
            "message": {
                "message_id": next(self._message_ids),
                "date": int(datetime.now(UTC).timestamp()),
                "chat": {"id": tg_id, "type": "private"},
                "from": {"id": 42, "is_bot": True, "first_name": "tgit"},
                "text": message_text,
                **({"reply_markup": reply_markup} if reply_markup else {}),
            },
        }
        await self._feed({"callback_query": callback})

    async def _feed(self, payload: dict[str, Any]) -> None:
        update = Update.model_validate({"update_id": next(self._update_ids), **payload}, context={"bot": self.bot})
        await self.dp.feed_update(self.bot, update)


def as_user(tg_id: int) -> User:
    return User(id=tg_id, is_bot=False, first_name=f"u{tg_id}")