diff --git a/README.md b/README.md index 6a7081f2..2a250281 100644 --- a/README.md +++ b/README.md @@ -203,6 +203,44 @@ app: FastAPI = create_fastapi_bot_app( - опционально `GET /metrics` - healthcheck routes через `setup_healthcheck(...)` +### Динамические account-ы и несколько CTS + +Если credentials появляются во время работы процесса, передайте +`account_provider` вместо `bot_accounts`. Provider является синхронным +read-side контрактом: он не должен выполнять сетевой или database I/O в +обработчике команды. Для динамического хранилища используйте локальный +snapshot, обновляемый приложением после регистрации или через invalidation. + +```python +from uuid import UUID + +from pybotx import Bot, BotAccountKey, BotAccountProvider, HandlerCollector + +collector = HandlerCollector() + + +def create_bot(registration_account_provider: BotAccountProvider) -> Bot: + return Bot( + collectors=[collector], + account_provider=registration_account_provider, + ) + +# В background-задаче CTS выбирается явно. +async def send_registration_confirmation( + bot: Bot, *, bot_id: UUID, chat_id: UUID +) -> None: + with bot.account_scope(BotAccountKey(server_id="cts-001", bot_id=bot_id)): + await bot.send_message(bot_id=bot_id, chat_id=chat_id, body="Готово") +``` + +Входящие команды связываются с `BotAccountKey` provider-ом до вызова handler; +JWT v1 cache изолирован ключом account-а и его `revision`. Если вызов производится +вне handler и один `bot_id` зарегистрирован на нескольких CTS, provider обязан +поднять `AmbiguousBotAccountError`, а не выбрать произвольный credential. + +Полный контракт и требования к adapter-у описаны в +[`docs/dynamic_accounts.md`](docs/dynamic_accounts.md). + Если у вас уже есть свой `FastAPI()` объект, можно не создавать новый app, а подключить `pybotx` в существующий: diff --git a/pybotx/__init__.py b/pybotx/__init__.py index 197f8ee5..eab27255 100644 --- a/pybotx/__init__.py +++ b/pybotx/__init__.py @@ -15,6 +15,10 @@ ) from pybotx.auth import BotXAuthVersion from pybotx.bot.bot import Bot +from pybotx.bot.bot_account_provider import ( + BotAccountProvider, + StaticBotAccountProvider, +) from pybotx.bot.command_processing import ( BotCommandOverloadAction, BotCommandOverloadStrategy, @@ -152,7 +156,11 @@ AttachmentVoice, OutgoingAttachment, ) -from pybotx.models.bot_account import BotAccount, BotAccountWithSecret +from pybotx.models.bot_account import ( + BotAccount, + BotAccountKey, + BotAccountWithSecret, +) from pybotx.models.bot_catalog import BotsListItem from pybotx.models.bot_sender import BotSender from pybotx.models.chats import ( @@ -242,6 +250,8 @@ "AttachmentVideo", "AttachmentVoice", "Bot", + "BotAccountKey", + "BotAccountProvider", "BotAPIBotDisabledErrorData", "BotAPIBotDisabledResponse", "BotAPIMethodFailedCallback", @@ -385,6 +395,7 @@ "SmartappManifestWebLayoutChoices", "SmartappManifestWebParams", "StatusRecipient", + "StaticBotAccountProvider", "StealthModeDisabledError", "Sticker", "StickerPack", diff --git a/pybotx/bot/bot.py b/pybotx/bot/bot.py index 4dd5f56d..e76751c6 100644 --- a/pybotx/bot/bot.py +++ b/pybotx/bot/bot.py @@ -1,6 +1,6 @@ from asyncio import Task from collections.abc import AsyncIterable, AsyncIterator, Iterator, Mapping, Sequence -from contextlib import AsyncExitStack, asynccontextmanager +from contextlib import AsyncExitStack, asynccontextmanager, contextmanager from contextvars import ContextVar, Token from datetime import datetime from types import SimpleNamespace @@ -14,6 +14,7 @@ from aiofiles.tempfile import NamedTemporaryFile, TemporaryDirectory from pybotx.async_buffer import AsyncBufferReadable, AsyncBufferWritable +from pybotx.bot.bot_account_provider import BotAccountProvider from pybotx.bot.bot_accounts_storage import BotAccountsStorage from pybotx.auth import BotXAuthVersion from pybotx.bot.callbacks.callback_manager import ( @@ -30,6 +31,7 @@ from pybotx.bot.callbacks.callback_memory_repo import CallbackMemoryRepo from pybotx.bot.callbacks.callback_repo_proto import CallbackRepoProto from pybotx.bot.contextvars import ( + bot_account_key_var, bot_id_var, chat_id_var, request_id_var, @@ -263,7 +265,7 @@ from pybotx.missing import Missing, MissingOptional, Undefined from pybotx.models.async_files import File from pybotx.models.attachments import IncomingFileAttachment, OutgoingAttachment -from pybotx.models.bot_account import BotAccountWithSecret +from pybotx.models.bot_account import BotAccountKey, BotAccountWithSecret from pybotx.models.bot_catalog import BotsListItem from pybotx.models.call import Call from pybotx.models.chats import ChatInfo, ChatLink, ChatListItem @@ -306,7 +308,8 @@ def __init__( self, *, collectors: Sequence[HandlerCollector], - bot_accounts: Sequence[BotAccountWithSecret], + bot_accounts: Sequence[BotAccountWithSecret] | None = None, + account_provider: BotAccountProvider | None = None, middlewares: Sequence[Middleware] | None = None, httpx_client: httpx.AsyncClient | None = None, httpx_timeout: httpx.Timeout | None = None, @@ -327,7 +330,11 @@ def __init__( ) -> None: if not collectors: logger.warning("Bot has no connected collectors") - if not bot_accounts: + if bot_accounts is not None and account_provider is not None: + raise ValueError("Pass either `bot_accounts` or `account_provider`, not both") + if bot_accounts is None and account_provider is None: + logger.warning("Bot has no bot accounts") + elif bot_accounts is not None and not bot_accounts: logger.warning("Bot has no bot accounts") middlewares = optional_sequence_to_list(middlewares) @@ -344,7 +351,8 @@ def __init__( self._default_callback_timeout = default_callback_timeout self._bot_accounts_storage = BotAccountsStorage( - list(bot_accounts), + bot_accounts, + account_provider=account_provider, auth_version=auth_version, retry_policy=retry_policy, retry_request_policy=retry_request_policy, @@ -444,10 +452,16 @@ def async_execute_bot_command( self, bot_command: BotCommand, ) -> "Task[None]": - # raise UnknownBotAccountError if no bot account with this bot_id. - self._bot_accounts_storage.ensure_bot_id_exists(bot_command.bot.id) + account_key = self._bot_accounts_storage.resolve_incoming_account_key( + bot_id=bot_command.bot.id, + host=bot_command.bot.host, + ) - return self._handler_collector.async_handle_bot_command(self, bot_command) + return self._handler_collector.async_handle_bot_command( + self, + bot_command, + account_key=account_key, + ) async def sync_execute_raw_smartapp_event( self, @@ -489,11 +503,18 @@ async def sync_execute_smartapp_event( self, smartapp_event: SmartAppEvent, ) -> BotAPISyncSmartAppEventResponse: - self._bot_accounts_storage.ensure_bot_id_exists(smartapp_event.bot.id) - return await self._handler_collector.handle_sync_smartapp_event( - self, - smartapp_event, + account_key = self._bot_accounts_storage.resolve_incoming_account_key( + bot_id=smartapp_event.bot.id, + host=smartapp_event.bot.host, ) + account_key_token = bot_account_key_var.set(account_key) + try: + return await self._handler_collector.handle_sync_smartapp_event( + self, + smartapp_event, + ) + finally: + bot_account_key_var.reset(account_key_token) async def raw_get_status( self, @@ -591,12 +612,32 @@ async def wait_botx_method_callback( def bot_accounts(self) -> Iterator[BotAccountWithSecret]: yield from self._bot_accounts_storage.iter_bot_accounts() + @contextmanager + def account_scope(self, account_key: BotAccountKey) -> Iterator[None]: + """Bind an explicit CTS identity for an outbound background operation.""" + self._bot_accounts_storage.ensure_account_key_exists(account_key) + account_key_token = bot_account_key_var.set(account_key) + try: + yield + finally: + bot_account_key_var.reset(account_key_token) + + def invalidate_account(self, account_key: BotAccountKey) -> None: + """Discard cached legacy JWTs after registration lifecycle changes.""" + self._bot_accounts_storage.invalidate(account_key) + async def fetch_tokens(self) -> None: if self._bot_accounts_storage.get_auth_version() != BotXAuthVersion.V1: return for bot_account in self.bot_accounts: try: - token = await self.get_token(bot_id=bot_account.id) + if bot_account.account_key is None: + token = await self.get_token(bot_id=bot_account.id) + self._bot_accounts_storage.set_token(bot_account.id, token) + else: + with self.account_scope(bot_account.account_key): + token = await self.get_token(bot_id=bot_account.id) + self._bot_accounts_storage.set_token(bot_account.id, token) except (InvalidBotAccountError, httpx.HTTPError): logger.opt(exception=True).warning( "Can't get token for bot account: " @@ -604,8 +645,6 @@ async def fetch_tokens(self) -> None: ) continue - self._bot_accounts_storage.set_token(bot_account.id, token) - async def startup(self, *, fetch_tokens: bool = True) -> None: if fetch_tokens: await self.fetch_tokens() diff --git a/pybotx/bot/bot_account_provider.py b/pybotx/bot/bot_account_provider.py new file mode 100644 index 00000000..5a86e777 --- /dev/null +++ b/pybotx/bot/bot_account_provider.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from typing import Protocol +from uuid import UUID + +from pybotx.bot.exceptions import UnknownBotAccountError +from pybotx.models.bot_account import BotAccountKey, BotAccountWithSecret + + +class BotAccountProvider(Protocol): + """Read-side port for resolving BotX credentials at runtime. + + Implementations must not perform blocking I/O. Dynamic providers should + serve a locally cached snapshot and refresh it outside a request path. + """ + + def get_account(self, key: BotAccountKey) -> BotAccountWithSecret: + """Return the credential for an explicit CTS and bot identity.""" + + def get_account_by_bot_id(self, bot_id: UUID) -> BotAccountWithSecret: + """Resolve only when ``bot_id`` is unambiguous for this provider.""" + + def resolve_incoming_account_key( + self, + *, + bot_id: UUID, + host: str | None, + ) -> BotAccountKey: + """Map inbound Bot API identity to a stable account key.""" + + def iter_bot_accounts(self) -> Iterator[BotAccountWithSecret]: + """Yield known accounts for legacy startup token fetching.""" + + +class StaticBotAccountProvider: + """Backward-compatible adapter over the historical ``bot_accounts`` list.""" + + def __init__(self, bot_accounts: Sequence[BotAccountWithSecret]) -> None: + self._bot_accounts = tuple(bot_accounts) + + def get_account(self, key: BotAccountKey) -> BotAccountWithSecret: + for account in self._bot_accounts: + if account.id == key.bot_id: + return account + raise UnknownBotAccountError(key.bot_id) + + def get_account_by_bot_id(self, bot_id: UUID) -> BotAccountWithSecret: + for account in self._bot_accounts: + if account.id == bot_id: + return account + raise UnknownBotAccountError(bot_id) + + def resolve_incoming_account_key( + self, + *, + bot_id: UUID, + host: str | None, + ) -> BotAccountKey: + account = self.get_account_by_bot_id(bot_id) + return BotAccountKey(server_id=host or account.host, bot_id=bot_id) + + def iter_bot_accounts(self) -> Iterator[BotAccountWithSecret]: + yield from self._bot_accounts diff --git a/pybotx/bot/bot_accounts_storage.py b/pybotx/bot/bot_accounts_storage.py index 2f960469..98a2d6e7 100644 --- a/pybotx/bot/bot_accounts_storage.py +++ b/pybotx/bot/bot_accounts_storage.py @@ -1,11 +1,15 @@ import base64 import hashlib import hmac -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from uuid import UUID from pybotx.auth import BotXAuthVersion, build_botx_jwt_v2 -from pybotx.bot.exceptions import UnknownBotAccountError +from pybotx.bot.bot_account_provider import ( + BotAccountProvider, + StaticBotAccountProvider, +) +from pybotx.bot.contextvars import bot_account_key_var from pybotx.client.http_config import ( BotXRetryPolicy, BotXRetryRequestPolicy, @@ -13,13 +17,15 @@ SafeBotXRetryRequestPolicy, ) from pybotx.client.observability import BotXRequestObserver -from pybotx.models.bot_account import BotAccountWithSecret +from pybotx.models.bot_account import BotAccountKey, BotAccountWithSecret class BotAccountsStorage: def __init__( self, - bot_accounts: list[BotAccountWithSecret], + bot_accounts: Sequence[BotAccountWithSecret] | None = None, + *, + account_provider: BotAccountProvider | None = None, auth_version: BotXAuthVersion = BotXAuthVersion.V2, retry_policy: BotXRetryPolicy | None = None, retry_request_policy: BotXRetryRequestPolicy | None = None, @@ -27,8 +33,14 @@ def __init__( metrics_collector: BotXRequestObserver | None = None, tracing_collector: BotXRequestObserver | None = None, ) -> None: - self._bot_accounts = bot_accounts - self._auth_tokens: dict[UUID, str] = {} + if bot_accounts is not None and account_provider is not None: + raise ValueError( + "Pass either `bot_accounts` or `account_provider`, not both" + ) + self._account_provider = account_provider or StaticBotAccountProvider( + bot_accounts or () + ) + self._auth_tokens: dict[tuple[BotAccountKey | UUID, int], str] = {} self._auth_version = auth_version self._retry_policy = retry_policy self._retry_request_policy = ( @@ -43,14 +55,27 @@ def __init__( self._tracing_collector = tracing_collector def get_bot_account(self, bot_id: UUID) -> BotAccountWithSecret: - for bot_account in self._bot_accounts: - if bot_account.id == bot_id: - return bot_account + account_key = self._get_current_account_key(bot_id) + if account_key is not None: + return self._account_provider.get_account(account_key) + return self._account_provider.get_account_by_bot_id(bot_id) + + def resolve_incoming_account_key( + self, + *, + bot_id: UUID, + host: str | None, + ) -> BotAccountKey: + return self._account_provider.resolve_incoming_account_key( + bot_id=bot_id, + host=host, + ) - raise UnknownBotAccountError(bot_id) + def ensure_account_key_exists(self, account_key: BotAccountKey) -> None: + self._account_provider.get_account(account_key) def iter_bot_accounts(self) -> Iterator[BotAccountWithSecret]: - yield from self._bot_accounts + yield from self._account_provider.iter_bot_accounts() def get_auth_version(self) -> BotXAuthVersion: return self._auth_version @@ -75,10 +100,10 @@ def get_cts_url(self, bot_id: UUID) -> str: return str(bot_account.cts_url) def set_token(self, bot_id: UUID, token: str) -> None: - self._auth_tokens[bot_id] = token + self._auth_tokens[self._token_cache_key(bot_id)] = token def get_token_or_none(self, bot_id: UUID) -> str | None: - return self._auth_tokens.get(bot_id) + return self._auth_tokens.get(self._token_cache_key(bot_id)) def build_jwt_v2(self, bot_id: UUID) -> str: bot_account = self.get_bot_account(bot_id) @@ -101,3 +126,23 @@ def build_signature(self, bot_id: UUID) -> str: def ensure_bot_id_exists(self, bot_id: UUID) -> None: self.get_bot_account(bot_id) + + def invalidate(self, account_key: BotAccountKey) -> None: + for cache_key in tuple(self._auth_tokens): + if cache_key[0] == account_key: + self._auth_tokens.pop(cache_key, None) + + def _token_cache_key(self, bot_id: UUID) -> tuple[BotAccountKey | UUID, int]: + account = self.get_bot_account(bot_id) + account_key = self._get_current_account_key(bot_id) + return (account_key or account.account_key or account.id, account.revision) + + @staticmethod + def _get_current_account_key(bot_id: UUID) -> BotAccountKey | None: + try: + account_key = bot_account_key_var.get() + except LookupError: + return None + if account_key.bot_id == bot_id: + return account_key + return None diff --git a/pybotx/bot/contextvars.py b/pybotx/bot/contextvars.py index 8b7595e9..37420c82 100644 --- a/pybotx/bot/contextvars.py +++ b/pybotx/bot/contextvars.py @@ -2,11 +2,14 @@ from typing import TYPE_CHECKING from uuid import UUID +from pybotx.models.bot_account import BotAccountKey + if TYPE_CHECKING: # To avoid circular import from pybotx.bot.bot import Bot bot_var: ContextVar["Bot"] = ContextVar("bot_var") bot_id_var: ContextVar[UUID] = ContextVar("bot_id") +bot_account_key_var: ContextVar[BotAccountKey] = ContextVar("bot_account_key") chat_id_var: ContextVar[UUID] = ContextVar("chat_id") request_id_var: ContextVar[str] = ContextVar("request_id") trace_id_var: ContextVar[str] = ContextVar("trace_id") diff --git a/pybotx/bot/exceptions.py b/pybotx/bot/exceptions.py index ace2f5f7..92436b1c 100644 --- a/pybotx/bot/exceptions.py +++ b/pybotx/bot/exceptions.py @@ -9,6 +9,15 @@ def __init__(self, bot_id: UUID) -> None: super().__init__(self.message) +class AmbiguousBotAccountError(Exception): + def __init__(self, bot_id: UUID) -> None: + self.bot_id = bot_id + super().__init__( + f"More than one bot account matches bot_id: `{bot_id!s}`. " + "Use an explicit BotAccountKey." + ) + + class BotXMethodCallbackNotFoundError(Exception): def __init__(self, sync_id: UUID) -> None: self.sync_id = sync_id diff --git a/pybotx/bot/handler_collector.py b/pybotx/bot/handler_collector.py index 92bb02a9..f64882f5 100644 --- a/pybotx/bot/handler_collector.py +++ b/pybotx/bot/handler_collector.py @@ -15,6 +15,7 @@ BotCommandProcessingConfig, ) from pybotx.bot.contextvars import ( + bot_account_key_var, bot_id_var, bot_var, chat_id_var, @@ -48,6 +49,7 @@ from pybotx.converters import optional_sequence_to_list from pybotx.logger import logger from pybotx.models.commands import BotCommand, SystemEvent +from pybotx.models.bot_account import BotAccountKey from pybotx.models.message.incoming_message import IncomingMessage from pybotx.models.status import BotMenu, StatusRecipient from pybotx.models.sync_smartapp_event import BotAPISyncSmartAppEventResponse @@ -86,9 +88,17 @@ class _QueuedBotCommand: request_id: str | None trace_id: str | None ingress_metadata: IngressCommandMetadata + account_key: BotAccountKey | None = None async def execute(self, collector: "HandlerCollector") -> None: - await collector.handle_bot_command(self.bot_command, self.bot) + account_key_token = None + if self.account_key is not None: + account_key_token = bot_account_key_var.set(self.account_key) + try: + await collector.handle_bot_command(self.bot_command, self.bot) + finally: + if account_key_token is not None: + bot_account_key_var.reset(account_key_token) class HandlerCollector: @@ -141,6 +151,8 @@ def async_handle_bot_command( self, bot: "Bot", bot_command: BotCommand, + *, + account_key: BotAccountKey | None = None, ) -> "asyncio.Task[None]": ingress_metadata = self._build_ingress_command_metadata(bot_command) request_id = self._get_optional_context_value(request_id_var) @@ -155,6 +167,7 @@ def async_handle_bot_command( request_id=request_id, trace_id=trace_id, ingress_metadata=ingress_metadata, + account_key=account_key, ) self._enqueue_or_reject(queued_command) diff --git a/pybotx/models/bot_account.py b/pybotx/models/bot_account.py index 66ca28de..1c976833 100644 --- a/pybotx/models/bot_account.py +++ b/pybotx/models/bot_account.py @@ -12,10 +12,20 @@ class BotAccount: host: str | None +@dataclass(frozen=True, slots=True) +class BotAccountKey: + """Stable account identity for one BotX bot in one CTS.""" + + server_id: str + bot_id: UUID + + class BotAccountWithSecret(BaseModel): id: UUID cts_url: AnyHttpUrl secret_key: str + account_key: BotAccountKey | None = None + revision: int = 1 model_config = ConfigDict(frozen=True) diff --git a/pyproject.toml b/pyproject.toml index 124a99ad..020863e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pybotx" -version = "0.77.0" +version = "0.77.1" description = "A python library for interacting with eXpress BotX API" authors = [ { name = "Sidnev Nikolay", email = "nsidnev@ccsteam.ru" }, diff --git a/tests/test_bot_account_provider.py b/tests/test_bot_account_provider.py new file mode 100644 index 00000000..e88acbeb --- /dev/null +++ b/tests/test_bot_account_provider.py @@ -0,0 +1,302 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Iterator +from uuid import UUID, uuid4 + +import pytest + +from pybotx import ( + Bot, + BotAccountKey, + BotAccountWithSecret, + BotXAuthVersion, + HandlerCollector, + StaticBotAccountProvider, + lifespan_wrapper, +) +from pybotx.bot.bot_accounts_storage import BotAccountsStorage +from pybotx.bot.contextvars import bot_account_key_var +from pybotx.bot.exceptions import AmbiguousBotAccountError, UnknownBotAccountError +from pybotx.models.bot_account import BotAccount +from pybotx.models.chats import Chat +from pybotx.models.enums import ChatTypes +from pybotx.models.message.incoming_message import ( + IncomingMessage, + UserDevice, + UserSender, +) + + +class SnapshotAccountProvider: + def __init__(self, accounts: list[BotAccountWithSecret]) -> None: + self._accounts: dict[BotAccountKey, BotAccountWithSecret] = {} + for account in accounts: + assert account.account_key is not None + self._accounts[account.account_key] = account + + def get_account(self, key: BotAccountKey) -> BotAccountWithSecret: + account = self._accounts.get(key) + if account is None: + raise UnknownBotAccountError(key.bot_id) + return account + + def get_account_by_bot_id(self, bot_id: UUID) -> BotAccountWithSecret: + matches = [account for account in self._accounts.values() if account.id == bot_id] + if not matches: + raise UnknownBotAccountError(bot_id) + if len(matches) > 1: + raise AmbiguousBotAccountError(bot_id) + return matches[0] + + def resolve_incoming_account_key( + self, *, bot_id: UUID, host: str | None + ) -> BotAccountKey: + matches = [ + key + for key, account in self._accounts.items() + if account.id == bot_id and account.host == host + ] + if len(matches) != 1: + raise UnknownBotAccountError(bot_id) + return matches[0] + + def iter_bot_accounts(self) -> Iterator[BotAccountWithSecret]: + yield from self._accounts.values() + + +@pytest.fixture +def dynamic_accounts(bot_id: UUID) -> tuple[BotAccountWithSecret, BotAccountWithSecret]: + first_key = BotAccountKey(server_id="cts-a", bot_id=bot_id) + second_key = BotAccountKey(server_id="cts-b", bot_id=bot_id) + return ( + BotAccountWithSecret( + id=bot_id, + cts_url="https://cts-a.example.com", + secret_key="secret-a", + account_key=first_key, + revision=1, + ), + BotAccountWithSecret( + id=bot_id, + cts_url="https://cts-b.example.com", + secret_key="secret-b", + account_key=second_key, + revision=2, + ), + ) + + +def test__accounts_storage__requires_explicit_context_for_ambiguous_bot_id( + bot_id: UUID, + dynamic_accounts: tuple[BotAccountWithSecret, BotAccountWithSecret], +) -> None: + storage = BotAccountsStorage(account_provider=SnapshotAccountProvider(list(dynamic_accounts))) + + with pytest.raises(AmbiguousBotAccountError): + storage.get_cts_url(bot_id) + + +def test__accounts_storage__validates_a_single_account_source( + bot_account: BotAccountWithSecret, +) -> None: + provider = StaticBotAccountProvider([bot_account]) + + with pytest.raises(ValueError, match="either `bot_accounts` or `account_provider`"): + Bot( + collectors=[HandlerCollector()], + bot_accounts=[bot_account], + account_provider=provider, + ) + + with pytest.raises(ValueError, match="either `bot_accounts` or `account_provider`"): + BotAccountsStorage([bot_account], account_provider=provider) + + +def test__accounts_storage__ignores_context_for_another_bot_id( + bot_id: UUID, + bot_account: BotAccountWithSecret, +) -> None: + storage = BotAccountsStorage([bot_account]) + context_token = bot_account_key_var.set( + BotAccountKey(server_id="another-cts", bot_id=uuid4()) + ) + try: + assert storage.get_cts_url(bot_id) == "https://cts.example.com/" + finally: + bot_account_key_var.reset(context_token) + + +def test__static_account_provider__raises_for_an_unknown_account_key( + bot_account: BotAccountWithSecret, +) -> None: + provider = StaticBotAccountProvider([bot_account]) + + with pytest.raises(UnknownBotAccountError): + provider.get_account(BotAccountKey(server_id="cts", bot_id=uuid4())) + + +def test__bot__warns_without_an_account_source( + loguru_caplog: pytest.LogCaptureFixture, +) -> None: + bot = Bot(collectors=[HandlerCollector()]) + try: + assert "Bot has no bot accounts" in loguru_caplog.text + finally: + asyncio.run(bot.shutdown()) + + +def test__accounts_storage__isolates_v1_tokens_by_account_key_and_revision( + bot_id: UUID, + dynamic_accounts: tuple[BotAccountWithSecret, BotAccountWithSecret], +) -> None: + first, second = dynamic_accounts + assert first.account_key is not None + assert second.account_key is not None + storage = BotAccountsStorage(account_provider=SnapshotAccountProvider([first, second])) + + first_token = bot_account_key_var.set(first.account_key) + try: + assert storage.get_cts_url(bot_id) == "https://cts-a.example.com/" + storage.set_token(bot_id, "token-a") + assert storage.get_token_or_none(bot_id) == "token-a" + finally: + bot_account_key_var.reset(first_token) + + second_token = bot_account_key_var.set(second.account_key) + try: + assert storage.get_cts_url(bot_id) == "https://cts-b.example.com/" + assert storage.get_token_or_none(bot_id) is None + storage.set_token(bot_id, "token-b") + finally: + bot_account_key_var.reset(second_token) + + storage.invalidate(first.account_key) + first_token = bot_account_key_var.set(first.account_key) + try: + assert storage.get_token_or_none(bot_id) is None + finally: + bot_account_key_var.reset(first_token) + second_token = bot_account_key_var.set(second.account_key) + try: + assert storage.get_token_or_none(bot_id) == "token-b" + finally: + bot_account_key_var.reset(second_token) + + +@pytest.mark.asyncio +async def test__bot__resolves_each_inbound_command_to_its_cts_context( + bot_id: UUID, + dynamic_accounts: tuple[BotAccountWithSecret, BotAccountWithSecret], +) -> None: + collector = HandlerCollector() + observed_urls: list[str] = [] + + @collector.command("/where", description="Resolve the current CTS") + async def where(message: IncomingMessage, bot: Bot) -> None: + await asyncio.sleep(0) + observed_urls.append(bot._bot_accounts_storage.get_cts_url(message.bot.id)) + + bot = Bot( + collectors=[collector], + account_provider=SnapshotAccountProvider(list(dynamic_accounts)), + ) + async with lifespan_wrapper(bot): + tasks = [] + for number, host in enumerate(("cts-a.example.com", "cts-b.example.com"), 1): + message = IncomingMessage( + bot=BotAccount(id=bot_id, host=host), + sync_id=UUID(int=number), + source_sync_id=None, + body="/where", + data={}, + metadata={}, + raw_command=None, + sender=UserSender( + huid=UUID(int=1), + udid=None, + ad_login=None, + ad_domain=None, + username=None, + is_chat_admin=None, + is_chat_creator=None, + device=UserDevice( + manufacturer=None, + device_name=None, + os=None, + pushes=None, + timezone=None, + permissions=None, + platform=None, + platform_package_id=None, + app_version=None, + locale=None, + ), + ), + chat=Chat(id=UUID(int=1), type=ChatTypes.GROUP_CHAT), + ) + tasks.append(bot.async_execute_bot_command(message)) + await asyncio.gather(*tasks) + + assert sorted(observed_urls) == [ + "https://cts-a.example.com/", + "https://cts-b.example.com/", + ] + + +def test__bot__account_scope_supports_background_operation_context( + bot_id: UUID, + dynamic_accounts: tuple[BotAccountWithSecret, BotAccountWithSecret], +) -> None: + first, second = dynamic_accounts + assert first.account_key is not None + assert second.account_key is not None + bot = Bot( + collectors=[HandlerCollector()], + account_provider=SnapshotAccountProvider([first, second]), + ) + + with bot.account_scope(first.account_key): + assert bot._bot_accounts_storage.get_cts_url(bot_id) == "https://cts-a.example.com/" + with bot.account_scope(second.account_key): + assert bot._bot_accounts_storage.get_cts_url(bot_id) == "https://cts-b.example.com/" + + with bot.account_scope(first.account_key): + bot._bot_accounts_storage.set_token(bot_id, "token-a") + bot.invalidate_account(first.account_key) + with bot.account_scope(first.account_key): + assert bot._bot_accounts_storage.get_token_or_none(bot_id) is None + + +@pytest.mark.asyncio +async def test__bot__fetch_tokens_binds_each_dynamic_account_scope( + bot_id: UUID, + dynamic_accounts: tuple[BotAccountWithSecret, BotAccountWithSecret], +) -> None: + first, second = dynamic_accounts + assert first.account_key is not None + assert second.account_key is not None + bot = Bot( + collectors=[HandlerCollector()], + account_provider=SnapshotAccountProvider([first, second]), + auth_version=BotXAuthVersion.V1, + ) + requested_urls: list[str] = [] + + async def get_token(*, bot_id: UUID) -> str: + requested_urls.append(bot._bot_accounts_storage.get_cts_url(bot_id)) + return f"token-{len(requested_urls)}" + + bot.get_token = get_token # type: ignore[method-assign] + try: + await bot.fetch_tokens() + assert requested_urls == [ + "https://cts-a.example.com/", + "https://cts-b.example.com/", + ] + with bot.account_scope(first.account_key): + assert bot._bot_accounts_storage.get_token_or_none(bot_id) == "token-1" + with bot.account_scope(second.account_key): + assert bot._bot_accounts_storage.get_token_or_none(bot_id) == "token-2" + finally: + await bot.shutdown()