Migrate motor to pymongo (#1705)

* migrated mongo storage from using deprecated motor to PyMongo

* added storages to __init__.py file to improve DX

* changelog file created

* Revert "added storages to __init__.py file to improve DX"

This reverts commit 5d0f6a9dfb.

* added optional dependency to pymongo to pyproject.toml

* Revert "migrated mongo storage from using deprecated motor to PyMongo"

This reverts commit 1c0207e1d1.

* added deprecation warning to mongo storage

* created pymongo storage

* added entry for PyMongoStorage to documentation in fsm.storages

* updated changelog to have information about how to migrate from MongoStorage to PyMongoStorage

* added test for pymongo storage (copied from mongo storage test)

* fixed formatting using black and isort

* fixed bug in close method of PyMongoStorage (client close method was not awaited)

* added test for PyMongoStorage that checks if storage could be properly closed

* pymongo package changed to be lower case in PyMongoStorage

* added fixture registration for pymongo storage

* test for pymongo is now using proper test fixtures

* removed redundant call to get_data, because we have checked this condition in the previous line

* added more tests to pymongo test, to check for all possible cases of using update_data method

* fixed PyMongoStorage update_data method implementation

* added pymongo tests to test_storages

* fixed pymongo tests, update_data method should not delete document when {} was passed

* Revert "fixed PyMongoStorage update_data method implementation"

This reverts commit 86170e1cb9.

* fixed linting issues in PyMongoStorage

* changed allowed versions of pymongo, to be compatible with motor

* pinned the upper version of pymongo to <4.11

tests/conftest.py

@@ -17,6 +17,7 @@ from aiogram.fsm.storage.memory import (
     SimpleEventIsolation,
 )
 from aiogram.fsm.storage.mongo import MongoStorage
+from aiogram.fsm.storage.pymongo import PyMongoStorage
 from aiogram.fsm.storage.redis import RedisStorage
 from tests.mocked_bot import MockedBot
 
@@ -102,6 +103,36 @@ async def mongo_storage(mongo_server):
         await storage.close()
 
 
+@pytest.fixture()
+def pymongo_server(request):
+    mongo_uri = request.config.getoption("--mongo")
+    if mongo_uri is None:
+        pytest.skip(SKIP_MESSAGE_PATTERN.format(db="mongo"))
+    else:
+        return mongo_uri
+
+
+@pytest.fixture()
+async def pymongo_storage(pymongo_server):
+    try:
+        parse_mongo_url(pymongo_server)
+    except InvalidURI as e:
+        raise UsageError(INVALID_URI_PATTERN.format(db="mongo", uri=pymongo_server, err=e))
+    storage = PyMongoStorage.from_url(
+        url=pymongo_server,
+        connection_kwargs={"serverSelectionTimeoutMS": 2000},
+    )
+    try:
+        await storage._client.server_info()
+    except PyMongoError as e:
+        pytest.fail(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/test_fsm/storage/test_pymongo.py

@@ -0,0 +1,218 @@
+import pytest
+from pymongo.errors import PyMongoError
+
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.pymongo import PyMongoStorage, StorageKey
+from tests.conftest import CHAT_ID, USER_ID
+
+PREFIX = "fsm"
+
+
+async def test_get_storage_passing_only_url(pymongo_server):
+    storage = PyMongoStorage.from_url(url=pymongo_server)
+    try:
+        await storage._client.server_info()
+    except PyMongoError as e:
+        pytest.fail(str(e))
+
+
+async def test_pymongo_storage_close_does_not_throw(pymongo_server):
+    storage = PyMongoStorage.from_url(url=pymongo_server)
+    try:
+        assert await storage.close() is None
+    except Exception as e:
+        pytest.fail(f"close() raised an exception: {e}")
+
+
+async def test_update_not_existing_data_with_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    assert await pymongo_storage.update_data(key=storage_key, data={}) == {}
+    assert await pymongo_storage._collection.find_one({}) is None
+
+
+async def test_update_not_existing_data_with_non_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    assert await pymongo_storage.update_data(key=storage_key, data={"key": "value"}) == {
+        "key": "value"
+    }
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "data": {"key": "value"},
+    }
+    await pymongo_storage._collection.delete_one({})
+
+
+async def test_update_existing_data_with_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    await pymongo_storage.set_data(key=storage_key, data={"key": "value"})
+    assert await pymongo_storage.update_data(key=storage_key, data={}) == {"key": "value"}
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "data": {"key": "value"},
+    }
+    await pymongo_storage._collection.delete_one({})
+
+
+async def test_update_existing_data_with_non_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    await pymongo_storage.set_data(key=storage_key, data={"key": "value"})
+    assert await pymongo_storage.update_data(key=storage_key, data={"key": "new_value"}) == {
+        "key": "new_value"
+    }
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "data": {"key": "new_value"},
+    }
+    await pymongo_storage._collection.delete_one({})
+
+
+async def test_document_life_cycle(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    await pymongo_storage.set_state(storage_key, "test")
+    await pymongo_storage.set_data(storage_key, {"key": "value"})
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "state": "test",
+        "data": {"key": "value"},
+    }
+    await pymongo_storage.set_state(storage_key, None)
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "data": {"key": "value"},
+    }
+    await pymongo_storage.set_data(storage_key, {})
+    assert await pymongo_storage._collection.find_one({}) is None
+
+
+class TestStateAndDataDoNotAffectEachOther:
+    async def test_state_and_data_do_not_affect_each_other_while_getting(
+        self,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        assert await pymongo_storage._collection.find_one({}) is None
+        await pymongo_storage.set_state(storage_key, "test")
+        await pymongo_storage.set_data(storage_key, {"key": "value"})
+        assert await pymongo_storage.get_state(storage_key) == "test"
+        assert await pymongo_storage.get_data(storage_key) == {"key": "value"}
+
+    async def test_data_do_not_affect_to_deleted_state_getting(
+        self,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        await pymongo_storage.set_state(storage_key, "test")
+        await pymongo_storage.set_data(storage_key, {"key": "value"})
+        await pymongo_storage.set_state(storage_key, None)
+        assert await pymongo_storage.get_state(storage_key) is None
+
+    async def test_state_do_not_affect_to_deleted_data_getting(
+        self,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        await pymongo_storage.set_state(storage_key, "test")
+        await pymongo_storage.set_data(storage_key, {"key": "value"})
+        await pymongo_storage.set_data(storage_key, {})
+        assert await pymongo_storage.get_data(storage_key) == {}
+
+    async def test_state_do_not_affect_to_updating_not_existing_data_with_empty_dictionary(
+        self,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        await pymongo_storage.set_state(storage_key, "test")
+        assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == {
+            "state": "test"
+        }
+        assert await pymongo_storage.update_data(key=storage_key, data={}) == {}
+        assert await pymongo_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,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        await pymongo_storage.set_state(storage_key, "test")
+        assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == {
+            "state": "test"
+        }
+        assert await pymongo_storage.update_data(
+            key=storage_key,
+            data={"key": "value"},
+        ) == {"key": "value"}
+        assert await pymongo_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,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        await pymongo_storage.set_state(storage_key, "test")
+        await pymongo_storage.set_data(storage_key, {"key": "value"})
+        assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == {
+            "state": "test",
+            "data": {"key": "value"},
+        }
+        assert await pymongo_storage.update_data(key=storage_key, data={}) == {"key": "value"}
+        assert await pymongo_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,
+        pymongo_storage: PyMongoStorage,
+        storage_key: StorageKey,
+    ):
+        await pymongo_storage.set_state(storage_key, "test")
+        await pymongo_storage.set_data(storage_key, {"key": "value"})
+        assert await pymongo_storage._collection.find_one({}, projection={"_id": 0}) == {
+            "state": "test",
+            "data": {"key": "value"},
+        }
+        assert await pymongo_storage.update_data(
+            key=storage_key,
+            data={"key": "VALUE", "key_2": "value_2"},
+        ) == {"key": "VALUE", "key_2": "value_2"}
+        assert await pymongo_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, pymongo_storage: PyMongoStorage):
+    assert pymongo_storage.resolve_state(value) == result

tests/test_fsm/storage/test_storages.py

@@ -11,6 +11,7 @@ from aiogram.fsm.storage.base import BaseStorage, StorageKey
     [
         pytest.lazy_fixture("redis_storage"),
         pytest.lazy_fixture("mongo_storage"),
+        pytest.lazy_fixture("pymongo_storage"),
         pytest.lazy_fixture("memory_storage"),
     ],
 )