tgit/tests/fakes.py

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

from collections.abc import AsyncGenerator
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.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"

    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, 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) -> None:
        self.dp = dp
        self.session = FakeSession()
        self.bot = Bot(BOT_TOKEN, session=self.session)
        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(self, tg_id: int, data: str, message_text: str = "section") -> 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,
            },
        }
        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}")