Added MongoStorage for FSM (#1434)

* Mongo storage included to storages test

* Added few additional checks in storages test

* Added MongoStorage for FSM

* Added changes description

* Fixed error message syntax

Co-authored-by: Alex Root Junior <jroot.junior@gmail.com>

* Resolved mypy check error

* IF/ELSE statement simplified

* Fix ruff linter error: RET505 Unnecessary `elif` after `return` statement

* Fix ruff linter error: E501 Line too long (100 > 99)

* Added mongo storage testing in CI

* Refactoring while review

* Refactoring while review

* Storing FSM state and data together in MongoDB-storage

* Fix CI - MongoDB container action is only supported on Linux

* Refactoring while review

* Enable Macos in pypy-tests section of CI

* Refactoring while review

* Makefile updated

* redis and mongo storages tests do not run in pypy-tests job of CI

* Fix docstring of DefaultKeyBuilder

---------

Co-authored-by: Alex Root Junior <jroot.junior@gmail.com>

.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: |

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)

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

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:`<prefix>:<bot_id?>:<business_connection_id?>:<chat_id>:<user_id>:<destiny?>:<field?>`
+    """
+
+    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

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", {})

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:`<prefix>:<bot_id?>:<business_connection_id?>:<chat_id>:<user_id>:<destiny?>:<field>`
-    """
-
-    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`)

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://<host>:<port>/<db> tests
+    pytest --redis redis://<host>:<port>/<db> --mongo mongodb://<user>:<password>@<host>:<port> tests
 
 Docs
 ----

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
 

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"
 
 

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()

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"

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)

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

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

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",
+        }