diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 181bf2ca..90afe538 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -44,8 +44,11 @@ jobs: env: # Windows has some limitations: - # – Redis is not supported on GitHub Windows runners; + # – Redis and MongoDB is not supported on GitHub Windows runners; IS_WINDOWS: ${{ startswith(matrix.os, 'windows') }} + # MongoDB has some limitations: + # – MongoDB container action is only supported on Linux; + IS_UBUNTU: ${{ startswith(matrix.os, 'ubuntu') }} steps: - name: Checkout code @@ -60,7 +63,7 @@ jobs: - name: Install project dependencies run: | - pip install -e .[dev,test,redis,proxy,i18n,fast] + pip install -e .[dev,test,redis,mongo,proxy,i18n,fast] - name: Lint code run: | @@ -74,10 +77,20 @@ jobs: with: redis-version: 6 + - name: Setup mongodb + if: ${{ env.IS_UBUNTU == 'true' }} + uses: supercharge/mongodb-github-action@1.10.0 + with: + mongodb-version: '7.0' + mongodb-username: mongo + mongodb-password: mongo + mongodb-port: 27017 + - name: Run tests run: | flags="$flags --cov=aiogram --cov-config .coveragerc --cov-report=xml" [[ "$IS_WINDOWS" == "false" ]] && flags="$flags --redis redis://localhost:6379/0" + [[ "$IS_UBUNTU" == "true" ]] && flags="$flags --mongo mongodb://mongo:mongo@localhost:27017" pytest $flags - name: Upload coverage data @@ -122,12 +135,7 @@ jobs: - name: Install project dependencies run: | - pip install -e .[dev,test,redis,proxy,i18n,fast] - - - name: Setup redis - uses: shogo82148/actions-setup-redis@v1 - with: - redis-version: 6 + pip install -e .[dev,test,redis,mongo,proxy,i18n,fast] - name: Run tests run: | diff --git a/CHANGES/1434.feature.rst b/CHANGES/1434.feature.rst new file mode 100644 index 00000000..1d6fcf6a --- /dev/null +++ b/CHANGES/1434.feature.rst @@ -0,0 +1 @@ +Added new storage :code:`aiogram.fsm.storage.MongoStorage` for Finite State Machine based on Mongo DB (using :code:`motor` library) diff --git a/Makefile b/Makefile index 128ff5fe..324d95c7 100644 --- a/Makefile +++ b/Makefile @@ -8,6 +8,7 @@ code_dir := $(package_dir) $(tests_dir) $(scripts_dir) $(examples_dir) reports_dir := reports redis_connection := redis://localhost:6379 +mongo_connection := mongodb://mongo:mongo@localhost:27017 # ================================================================================================= # Environment @@ -50,12 +51,12 @@ test-run-services: .PHONY: test test: test-run-services - pytest --cov=aiogram --cov-config .coveragerc tests/ --redis $(redis_connection) + pytest --cov=aiogram --cov-config .coveragerc tests/ --redis $(redis_connection) --mongo $(mongo_connection) .PHONY: test-coverage test-coverage: test-run-services mkdir -p $(reports_dir)/tests/ - pytest --cov=aiogram --cov-config .coveragerc --html=$(reports_dir)/tests/index.html tests/ --redis $(redis_connection) + pytest --cov=aiogram --cov-config .coveragerc --html=$(reports_dir)/tests/index.html tests/ --redis $(redis_connection) --mongo $(mongo_connection) coverage html -d $(reports_dir)/coverage .PHONY: test-coverage-view diff --git a/aiogram/fsm/storage/base.py b/aiogram/fsm/storage/base.py index a66d56be..8e1b206f 100644 --- a/aiogram/fsm/storage/base.py +++ b/aiogram/fsm/storage/base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, AsyncGenerator, Dict, Optional, Union +from typing import Any, AsyncGenerator, Dict, Literal, Optional, Union from aiogram.fsm.state import State @@ -20,6 +20,85 @@ class StorageKey: destiny: str = DEFAULT_DESTINY +class KeyBuilder(ABC): + """Base class for key builder.""" + + @abstractmethod + def build( + self, + key: StorageKey, + part: Optional[Literal["data", "state", "lock"]] = None, + ) -> str: + """ + Build key to be used in storage's db queries + + :param key: contextual key + :param part: part of the record + :return: key to be used in storage's db queries + """ + pass + + +class DefaultKeyBuilder(KeyBuilder): + """ + Simple key builder with default prefix. + + Generates a colon-joined string with prefix, chat_id, user_id, + optional bot_id, business_connection_id, destiny and field. + + Format: + :code:`::::::` + """ + + def __init__( + self, + *, + prefix: str = "fsm", + separator: str = ":", + with_bot_id: bool = False, + with_business_connection_id: bool = False, + with_destiny: bool = False, + ) -> None: + """ + :param prefix: prefix for all records + :param separator: separator + :param with_bot_id: include Bot id in the key + :param with_business_connection_id: include business connection id + :param with_destiny: include destiny key + """ + self.prefix = prefix + self.separator = separator + self.with_bot_id = with_bot_id + self.with_business_connection_id = with_business_connection_id + self.with_destiny = with_destiny + + def build( + self, + key: StorageKey, + part: Optional[Literal["data", "state", "lock"]] = None, + ) -> str: + parts = [self.prefix] + if self.with_bot_id: + parts.append(str(key.bot_id)) + if self.with_business_connection_id and key.business_connection_id: + parts.append(str(key.business_connection_id)) + parts.append(str(key.chat_id)) + if key.thread_id: + parts.append(str(key.thread_id)) + parts.append(str(key.user_id)) + if self.with_destiny: + parts.append(key.destiny) + elif key.destiny != DEFAULT_DESTINY: + error_message = ( + "Default key builder is not configured to use key destiny other than the default." + "\n\nProbably, you should set `with_destiny=True` in for DefaultKeyBuilder." + ) + raise ValueError(error_message) + if part: + parts.append(part) + return self.separator.join(parts) + + class BaseStorage(ABC): """ Base class for all FSM storages diff --git a/aiogram/fsm/storage/mongo.py b/aiogram/fsm/storage/mongo.py new file mode 100644 index 00000000..b4b1eeaa --- /dev/null +++ b/aiogram/fsm/storage/mongo.py @@ -0,0 +1,130 @@ +from typing import Any, Dict, Optional, cast + +from motor.motor_asyncio import AsyncIOMotorClient + +from aiogram.fsm.state import State +from aiogram.fsm.storage.base import ( + BaseStorage, + DefaultKeyBuilder, + KeyBuilder, + StateType, + StorageKey, +) + + +class MongoStorage(BaseStorage): + """ + MongoDB storage required :code:`motor` package installed (:code:`pip install motor`) + """ + + def __init__( + self, + client: AsyncIOMotorClient, + key_builder: Optional[KeyBuilder] = None, + db_name: str = "aiogram_fsm", + collection_name: str = "states_and_data", + ) -> None: + """ + :param client: Instance of AsyncIOMotorClient + :param key_builder: builder that helps to convert contextual key to string + :param db_name: name of the MongoDB database for FSM + :param collection_name: name of the collection for storing FSM states and data + """ + if key_builder is None: + key_builder = DefaultKeyBuilder() + self._client = client + self._database = self._client[db_name] + self._collection = self._database[collection_name] + self._key_builder = key_builder + + @classmethod + def from_url( + cls, url: str, connection_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Any + ) -> "MongoStorage": + """ + Create an instance of :class:`MongoStorage` with specifying the connection string + + :param url: for example :code:`mongodb://user:password@host:port` + :param connection_kwargs: see :code:`motor` docs + :param kwargs: arguments to be passed to :class:`MongoStorage` + :return: an instance of :class:`MongoStorage` + """ + if connection_kwargs is None: + connection_kwargs = {} + client = AsyncIOMotorClient(url, **connection_kwargs) + return cls(client=client, **kwargs) + + async def close(self) -> None: + """Cleanup client resources and disconnect from MongoDB.""" + self._client.close() + + def resolve_state(self, value: StateType) -> Optional[str]: + if value is None: + return None + if isinstance(value, State): + return value.state + return str(value) + + async def set_state(self, key: StorageKey, state: StateType = None) -> None: + document_id = self._key_builder.build(key) + if state is None: + updated = await self._collection.find_one_and_update( + filter={"_id": document_id}, + update={"$unset": {"state": 1}}, + projection={"_id": 0}, + return_document=True, + ) + if updated == {}: + await self._collection.delete_one({"_id": document_id}) + else: + await self._collection.update_one( + filter={"_id": document_id}, + update={"$set": {"state": self.resolve_state(state)}}, + upsert=True, + ) + + async def get_state(self, key: StorageKey) -> Optional[str]: + document_id = self._key_builder.build(key) + document = await self._collection.find_one({"_id": document_id}) + if document is None: + return None + return document.get("state") + + async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None: + document_id = self._key_builder.build(key) + if not data: + updated = await self._collection.find_one_and_update( + filter={"_id": document_id}, + update={"$unset": {"data": 1}}, + projection={"_id": 0}, + return_document=True, + ) + if updated == {}: + await self._collection.delete_one({"_id": document_id}) + else: + await self._collection.update_one( + filter={"_id": document_id}, + update={"$set": {"data": data}}, + upsert=True, + ) + + async def get_data(self, key: StorageKey) -> Dict[str, Any]: + document_id = self._key_builder.build(key) + document = await self._collection.find_one({"_id": document_id}) + if document is None or not document.get("data"): + return {} + return cast(Dict[str, Any], document["data"]) + + async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]: + document_id = self._key_builder.build(key) + update_with = {f"data.{key}": value for key, value in data.items()} + update_result = await self._collection.find_one_and_update( + filter={"_id": document_id}, + update={"$set": update_with}, + upsert=True, + return_document=True, + projection={"_id": 0}, + ) + if not update_result: + await self._collection.delete_one({"_id": document_id}) + return update_result.get("data", {}) diff --git a/aiogram/fsm/storage/redis.py b/aiogram/fsm/storage/redis.py index a818a831..edcf9d63 100644 --- a/aiogram/fsm/storage/redis.py +++ b/aiogram/fsm/storage/redis.py @@ -1,7 +1,6 @@ import json -from abc import ABC, abstractmethod from contextlib import asynccontextmanager -from typing import Any, AsyncGenerator, Callable, Dict, Literal, Optional, cast +from typing import Any, AsyncGenerator, Callable, Dict, Optional, cast from redis.asyncio.client import Redis from redis.asyncio.connection import ConnectionPool @@ -10,9 +9,10 @@ from redis.typing import ExpiryT from aiogram.fsm.state import State from aiogram.fsm.storage.base import ( - DEFAULT_DESTINY, BaseEventIsolation, BaseStorage, + DefaultKeyBuilder, + KeyBuilder, StateType, StorageKey, ) @@ -22,79 +22,6 @@ _JsonLoads = Callable[..., Any] _JsonDumps = Callable[..., str] -class KeyBuilder(ABC): - """ - Base class for Redis key builder - """ - - @abstractmethod - def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str: - """ - This method should be implemented in subclasses - - :param key: contextual key - :param part: part of the record - :return: key to be used in Redis queries - """ - pass - - -class DefaultKeyBuilder(KeyBuilder): - """ - Simple Redis key builder with default prefix. - - Generates a colon-joined string with prefix, chat_id, user_id, - optional bot_id, business_connection_id and destiny. - - Format: - :code:`::::::` - """ - - def __init__( - self, - *, - prefix: str = "fsm", - separator: str = ":", - with_bot_id: bool = False, - with_business_connection_id: bool = False, - with_destiny: bool = False, - ) -> None: - """ - :param prefix: prefix for all records - :param separator: separator - :param with_bot_id: include Bot id in the key - :param with_business_connection_id: include business connection id - :param with_destiny: include a destiny key - """ - self.prefix = prefix - self.separator = separator - self.with_bot_id = with_bot_id - self.with_business_connection_id = with_business_connection_id - self.with_destiny = with_destiny - - def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str: - parts = [self.prefix] - if self.with_bot_id: - parts.append(str(key.bot_id)) - if self.with_business_connection_id and key.business_connection_id: - parts.append(str(key.business_connection_id)) - parts.append(str(key.chat_id)) - if key.thread_id: - parts.append(str(key.thread_id)) - parts.append(str(key.user_id)) - if self.with_destiny: - parts.append(key.destiny) - elif key.destiny != DEFAULT_DESTINY: - raise ValueError( - "Redis key builder is not configured to use key destiny other the default.\n" - "\n" - "Probably, you should set `with_destiny=True` in for DefaultKeyBuilder.\n" - "E.g: `RedisStorage(redis, key_builder=DefaultKeyBuilder(with_destiny=True))`" - ) - parts.append(part) - return self.separator.join(parts) - - class RedisStorage(BaseStorage): """ Redis storage required :code:`redis` package installed (:code:`pip install redis`) diff --git a/docs/contributing.rst b/docs/contributing.rst index a4a3b8e0..c68480b0 100644 --- a/docs/contributing.rst +++ b/docs/contributing.rst @@ -78,13 +78,13 @@ Linux / macOS: .. code-block:: bash - pip install -e ."[dev,test,docs,fast,redis,proxy,i18n]" + pip install -e ."[dev,test,docs,fast,redis,mongo,proxy,i18n]" Windows: .. code-block:: bash - pip install -e .[dev,test,docs,fast,redis,proxy,i18n] + pip install -e .[dev,test,docs,fast,redis,mongo,proxy,i18n] It will install :code:`aiogram` in editable mode into your virtual environment and all dependencies. @@ -116,11 +116,12 @@ All changes should be tested: pytest tests -Also if you are doing something with Redis-storage, you will need to test everything works with Redis: +Also if you are doing something with Redis-storage or/and MongoDB-storage, +you will need to test everything works with Redis or/and MongoDB: .. code-block:: bash - pytest --redis redis://:/ tests + pytest --redis redis://:/ --mongo mongodb://:@: tests Docs ---- diff --git a/docs/dispatcher/finite_state_machine/storages.rst b/docs/dispatcher/finite_state_machine/storages.rst index 702838e5..3795d2f1 100644 --- a/docs/dispatcher/finite_state_machine/storages.rst +++ b/docs/dispatcher/finite_state_machine/storages.rst @@ -19,13 +19,23 @@ RedisStorage :members: __init__, from_url :member-order: bysource -Keys inside storage can be customized via key builders: +MongoStorage +------------ -.. autoclass:: aiogram.fsm.storage.redis.KeyBuilder +.. autoclass:: aiogram.fsm.storage.mongo.MongoStorage + :members: __init__, from_url + :member-order: bysource + +KeyBuilder +------------ + +Keys inside Redis and Mongo storages can be customized via key builders: + +.. autoclass:: aiogram.fsm.storage.base.KeyBuilder :members: :member-order: bysource -.. autoclass:: aiogram.fsm.storage.redis.DefaultKeyBuilder +.. autoclass:: aiogram.fsm.storage.base.DefaultKeyBuilder :members: :member-order: bysource diff --git a/pyproject.toml b/pyproject.toml index c4a7664a..131f5d35 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,9 @@ fast = [ redis = [ "redis[hiredis]~=5.0.1", ] +mongo = [ + "motor~=3.3.2", +] proxy = [ "aiohttp-socks~=0.8.3", ] @@ -105,6 +108,7 @@ dev = [ "toml~=0.10.2", "pre-commit~=3.5.0", "packaging~=23.1", + "motor-types~=1.0.0b4", ] [project.urls] @@ -117,6 +121,7 @@ features = [ "dev", "fast", "redis", + "mongo", "proxy", "i18n", "cli", @@ -136,6 +141,7 @@ lint = "ruff aiogram" features = [ "fast", "redis", + "mongo", "proxy", "i18n", "docs", @@ -150,6 +156,7 @@ features = [ "dev", "fast", "redis", + "mongo", "proxy", "i18n", "test", @@ -167,6 +174,7 @@ update = [ features = [ "fast", "redis", + "mongo", "proxy", "i18n", "test", @@ -182,6 +190,10 @@ cov-redis = [ "pytest --cov-config pyproject.toml --cov=aiogram --html=reports/py{matrix:python}/tests/index.html --redis {env:REDIS_DNS:'redis://localhost:6379'} {args}", "coverage html -d reports/py{matrix:python}/coverage", ] +cov-mongo = [ + "pytest --cov-config pyproject.toml --cov=aiogram --html=reports/py{matrix:python}/tests/index.html --mongo {env:MONGO_DNS:'mongodb://mongo:mongo@localhost:27017'} {args}", + "coverage html -d reports/py{matrix:python}/coverage", +] view-cov = "google-chrome-stable reports/py{matrix:python}/coverage/index.html" diff --git a/tests/conftest.py b/tests/conftest.py index 3e063ff1..f13fec10 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,8 @@ from pathlib import Path import pytest from _pytest.config import UsageError +from pymongo.errors import InvalidURI, PyMongoError +from pymongo.uri_parser import parse_uri as parse_mongo_url from redis.asyncio.connection import parse_url as parse_redis_url from aiogram import Dispatcher @@ -10,6 +12,7 @@ from aiogram.fsm.storage.memory import ( MemoryStorage, SimpleEventIsolation, ) +from aiogram.fsm.storage.mongo import MongoStorage from aiogram.fsm.storage.redis import RedisEventIsolation, RedisStorage from tests.mocked_bot import MockedBot @@ -18,24 +21,27 @@ DATA_DIR = Path(__file__).parent / "data" def pytest_addoption(parser): parser.addoption("--redis", default=None, help="run tests which require redis connection") + parser.addoption("--mongo", default=None, help="run tests which require mongo connection") def pytest_configure(config): config.addinivalue_line("markers", "redis: marked tests require redis connection to run") + config.addinivalue_line("markers", "mongo: marked tests require mongo connection to run") def pytest_collection_modifyitems(config, items): - redis_uri = config.getoption("--redis") - if redis_uri is None: - skip_redis = pytest.mark.skip(reason="need --redis option with redis URI to run") - for item in items: - if "redis" in item.keywords: - item.add_marker(skip_redis) - return - try: - parse_redis_url(redis_uri) - except ValueError as e: - raise UsageError(f"Invalid redis URI {redis_uri!r}: {e}") + for db, parse_uri in [("redis", parse_redis_url), ("mongo", parse_mongo_url)]: + uri = config.getoption(f"--{db}") + if uri is None: + skip = pytest.mark.skip(reason=f"need --{db} option with {db} URI to run") + for item in items: + if db in item.keywords: + item.add_marker(skip) + else: + try: + parse_uri(uri) + except (ValueError, InvalidURI) as e: + raise UsageError(f"Invalid {db} URI {uri!r}: {e}") @pytest.fixture() @@ -62,6 +68,29 @@ async def redis_storage(redis_server): await storage.close() +@pytest.fixture() +def mongo_server(request): + mongo_uri = request.config.getoption("--mongo") + return mongo_uri + + +@pytest.fixture() +@pytest.mark.mongo +async def mongo_storage(mongo_server): + if not mongo_server: + pytest.skip("MongoDB is not available here") + storage = MongoStorage.from_url(mongo_server) + try: + await storage._client.server_info() + except PyMongoError as e: + pytest.skip(str(e)) + else: + yield storage + await storage._client.drop_database(storage._database) + finally: + await storage.close() + + @pytest.fixture() async def memory_storage(): storage = MemoryStorage() diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index 453f5e5a..895908ba 100644 --- a/tests/docker-compose.yml +++ b/tests/docker-compose.yml @@ -5,3 +5,11 @@ services: image: redis:6-alpine ports: - "${REDIS_PORT-6379}:6379" + + mongo: + image: mongo:7.0.6 + environment: + MONGO_INITDB_ROOT_USERNAME: mongo + MONGO_INITDB_ROOT_PASSWORD: mongo + ports: + - "${MONGODB_PORT-27017}:27017" diff --git a/tests/test_fsm/storage/test_isolation.py b/tests/test_fsm/storage/test_isolation.py index 1fb21d55..c95a1cba 100644 --- a/tests/test_fsm/storage/test_isolation.py +++ b/tests/test_fsm/storage/test_isolation.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch import pytest from aiogram.fsm.storage.base import BaseEventIsolation, StorageKey -from aiogram.fsm.storage.redis import RedisEventIsolation +from aiogram.fsm.storage.redis import RedisEventIsolation, RedisStorage from tests.mocked_bot import MockedBot @@ -32,6 +32,14 @@ class TestIsolations: class TestRedisEventIsolation: + def test_create_isolation(self): + fake_redis = object() + storage = RedisStorage(redis=fake_redis) + isolation = storage.create_isolation() + assert isinstance(isolation, RedisEventIsolation) + assert isolation.redis is fake_redis + assert isolation.key_builder is storage.key_builder + def test_init_without_key_builder(self): redis = AsyncMock() isolation = RedisEventIsolation(redis=redis) diff --git a/tests/test_fsm/storage/test_redis.py b/tests/test_fsm/storage/test_key_builder.py similarity index 75% rename from tests/test_fsm/storage/test_redis.py rename to tests/test_fsm/storage/test_key_builder.py index 18c143e3..f62ac505 100644 --- a/tests/test_fsm/storage/test_redis.py +++ b/tests/test_fsm/storage/test_key_builder.py @@ -1,11 +1,8 @@ +from typing import Literal, Optional + import pytest -from aiogram.fsm.storage.base import DEFAULT_DESTINY, StorageKey -from aiogram.fsm.storage.redis import ( - DefaultKeyBuilder, - RedisEventIsolation, - RedisStorage, -) +from aiogram.fsm.storage.base import DEFAULT_DESTINY, DefaultKeyBuilder, StorageKey PREFIX = "test" BOT_ID = 42 @@ -16,9 +13,9 @@ BUSINESS_CONNECTION_ID = "4" FIELD = "data" -class TestRedisDefaultKeyBuilder: +class TestDefaultKeyBuilder: @pytest.mark.parametrize( - "key_builder,result", + "key_builder,field,result", [ [ DefaultKeyBuilder( @@ -27,40 +24,52 @@ class TestRedisDefaultKeyBuilder: with_destiny=True, with_business_connection_id=True, ), + FIELD, f"{PREFIX}:{BOT_ID}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}", ], [ DefaultKeyBuilder(prefix=PREFIX, with_bot_id=True, with_destiny=True), - f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}", + None, + f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}", ], [ DefaultKeyBuilder( prefix=PREFIX, with_bot_id=True, with_business_connection_id=True ), + FIELD, f"{PREFIX}:{BOT_ID}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{FIELD}", ], [ DefaultKeyBuilder(prefix=PREFIX, with_bot_id=True), - f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}", + None, + f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}", ], [ DefaultKeyBuilder( prefix=PREFIX, with_destiny=True, with_business_connection_id=True ), + FIELD, f"{PREFIX}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}", ], [ DefaultKeyBuilder(prefix=PREFIX, with_destiny=True), - f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}", + None, + f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}", ], [ DefaultKeyBuilder(prefix=PREFIX, with_business_connection_id=True), + FIELD, f"{PREFIX}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{FIELD}", ], - [DefaultKeyBuilder(prefix=PREFIX), f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"], + [DefaultKeyBuilder(prefix=PREFIX), None, f"{PREFIX}:{CHAT_ID}:{USER_ID}"], ], ) - async def test_generate_key(self, key_builder: DefaultKeyBuilder, result: str): + async def test_generate_key( + self, + key_builder: DefaultKeyBuilder, + field: Optional[Literal["data", "state", "lock"]], + result: str, + ): key = StorageKey( chat_id=CHAT_ID, user_id=USER_ID, @@ -68,7 +77,7 @@ class TestRedisDefaultKeyBuilder: business_connection_id=BUSINESS_CONNECTION_ID, destiny=DEFAULT_DESTINY, ) - assert key_builder.build(key, FIELD) == result + assert key_builder.build(key, field) == result async def test_destiny_check(self): key_builder = DefaultKeyBuilder( @@ -95,11 +104,3 @@ class TestRedisDefaultKeyBuilder: destiny=DEFAULT_DESTINY, ) assert key_builder.build(key, FIELD) == f"{PREFIX}:{CHAT_ID}:{THREAD_ID}:{USER_ID}:{FIELD}" - - def test_create_isolation(self): - fake_redis = object() - storage = RedisStorage(redis=fake_redis) - isolation = storage.create_isolation() - assert isinstance(isolation, RedisEventIsolation) - assert isolation.redis is fake_redis - assert isolation.key_builder is storage.key_builder diff --git a/tests/test_fsm/storage/test_mongo.py b/tests/test_fsm/storage/test_mongo.py new file mode 100644 index 00000000..48b95719 --- /dev/null +++ b/tests/test_fsm/storage/test_mongo.py @@ -0,0 +1,164 @@ +import pytest + +from aiogram.fsm.state import State +from aiogram.fsm.storage.mongo import MongoStorage, StorageKey +from tests.mocked_bot import MockedBot + +PREFIX = "fsm" +CHAT_ID = -42 +USER_ID = 42 + + +@pytest.fixture(name="storage_key") +def create_storage_key(bot: MockedBot): + return StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=bot.id) + + +async def test_update_not_existing_data_with_empty_dictionary( + mongo_storage: MongoStorage, + storage_key: StorageKey, +): + assert await mongo_storage._collection.find_one({}) is None + assert await mongo_storage.get_data(key=storage_key) == {} + assert await mongo_storage.update_data(key=storage_key, data={}) == {} + assert await mongo_storage._collection.find_one({}) is None + + +async def test_document_life_cycle( + mongo_storage: MongoStorage, + storage_key: StorageKey, +): + assert await mongo_storage._collection.find_one({}) is None + await mongo_storage.set_state(storage_key, "test") + await mongo_storage.set_data(storage_key, {"key": "value"}) + assert await mongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "state": "test", + "data": {"key": "value"}, + } + await mongo_storage.set_state(storage_key, None) + assert await mongo_storage._collection.find_one({}) == { + "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}", + "data": {"key": "value"}, + } + await mongo_storage.set_data(storage_key, {}) + assert await mongo_storage._collection.find_one({}) is None + + +class TestStateAndDataDoNotAffectEachOther: + async def test_state_and_data_do_not_affect_each_other_while_getting( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + assert await mongo_storage._collection.find_one({}) is None + await mongo_storage.set_state(storage_key, "test") + await mongo_storage.set_data(storage_key, {"key": "value"}) + assert await mongo_storage.get_state(storage_key) == "test" + assert await mongo_storage.get_data(storage_key) == {"key": "value"} + + async def test_data_do_not_affect_to_deleted_state_getting( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + await mongo_storage.set_data(storage_key, {"key": "value"}) + await mongo_storage.set_state(storage_key, None) + assert await mongo_storage.get_state(storage_key) is None + + async def test_state_do_not_affect_to_deleted_data_getting( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + await mongo_storage.set_data(storage_key, {"key": "value"}) + await mongo_storage.set_data(storage_key, {}) + assert await mongo_storage.get_data(storage_key) == {} + + async def test_state_do_not_affect_to_updating_not_existing_data_with_empty_dictionary( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test" + } + assert await mongo_storage.update_data(key=storage_key, data={}) == {} + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test" + } + + async def test_state_do_not_affect_to_updating_not_existing_data_with_non_empty_dictionary( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test" + } + assert await mongo_storage.update_data( + key=storage_key, + data={"key": "value"}, + ) == {"key": "value"} + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + + async def test_state_do_not_affect_to_updating_existing_data_with_empty_dictionary( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + await mongo_storage.set_data(storage_key, {"key": "value"}) + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + assert await mongo_storage.update_data(key=storage_key, data={}) == {"key": "value"} + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + + async def test_state_do_not_affect_to_updating_existing_data_with_non_empty_dictionary( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + await mongo_storage.set_data(storage_key, {"key": "value"}) + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "value"}, + } + assert await mongo_storage.update_data( + key=storage_key, + data={"key": "VALUE", "key_2": "value_2"}, + ) == {"key": "VALUE", "key_2": "value_2"} + assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == { + "state": "test", + "data": {"key": "VALUE", "key_2": "value_2"}, + } + + +@pytest.mark.parametrize( + "value,result", + [ + [None, None], + ["", ""], + ["text", "text"], + [State(), None], + [State(state="*"), "*"], + [State("text"), "@:text"], + [State("test", group_name="Test"), "Test:test"], + [[1, 2, 3], "[1, 2, 3]"], + ], +) +def test_resolve_state(value, result, mongo_storage: MongoStorage): + assert mongo_storage.resolve_state(value) == result diff --git a/tests/test_fsm/storage/test_storages.py b/tests/test_fsm/storage/test_storages.py index 3558e33d..1c1f87a2 100644 --- a/tests/test_fsm/storage/test_storages.py +++ b/tests/test_fsm/storage/test_storages.py @@ -11,7 +11,11 @@ def create_storage_key(bot: MockedBot): @pytest.mark.parametrize( "storage", - [pytest.lazy_fixture("redis_storage"), pytest.lazy_fixture("memory_storage")], + [ + pytest.lazy_fixture("redis_storage"), + pytest.lazy_fixture("mongo_storage"), + pytest.lazy_fixture("memory_storage"), + ], ) class TestStorages: async def test_set_state(self, bot: MockedBot, storage: BaseStorage, storage_key: StorageKey): @@ -35,6 +39,8 @@ class TestStorages: ): assert await storage.get_data(key=storage_key) == {} assert await storage.update_data(key=storage_key, data={"foo": "bar"}) == {"foo": "bar"} + assert await storage.update_data(key=storage_key, data={}) == {"foo": "bar"} + assert await storage.get_data(key=storage_key) == {"foo": "bar"} assert await storage.update_data(key=storage_key, data={"baz": "spam"}) == { "foo": "bar", "baz": "spam", @@ -43,3 +49,11 @@ class TestStorages: "foo": "bar", "baz": "spam", } + assert await storage.update_data(key=storage_key, data={"baz": "test"}) == { + "foo": "bar", + "baz": "test", + } + assert await storage.get_data(key=storage_key) == { + "foo": "bar", + "baz": "test", + }