diff --git a/aiogram/fsm/storage/base.py b/aiogram/fsm/storage/base.py index 2309b10a..fab5a6ca 100644 --- a/aiogram/fsm/storage/base.py +++ b/aiogram/fsm/storage/base.py @@ -23,7 +23,11 @@ class KeyBuilder(ABC): """Base class for key builder.""" @abstractmethod - def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str: + def build( + self, + key: StorageKey, + part: Optional[Literal["data", "state", "lock"]] = None, + ) -> str: """ Build key to be used in storage's db queries @@ -61,7 +65,11 @@ class DefaultKeyBuilder(KeyBuilder): self.with_bot_id = with_bot_id self.with_destiny = with_destiny - def build(self, key: StorageKey, part: Literal["data", "state", "lock"]) -> str: + 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)) @@ -77,7 +85,8 @@ class DefaultKeyBuilder(KeyBuilder): "\n\nProbably, you should set `with_destiny=True` in for DefaultKeyBuilder." ) raise ValueError(error_message) - parts.append(part) + if part: + parts.append(part) return self.separator.join(parts) diff --git a/aiogram/fsm/storage/mongo.py b/aiogram/fsm/storage/mongo.py index 8bc95217..ca0346ac 100644 --- a/aiogram/fsm/storage/mongo.py +++ b/aiogram/fsm/storage/mongo.py @@ -22,24 +22,19 @@ class MongoStorage(BaseStorage): client: AsyncIOMotorClient, key_builder: Optional[KeyBuilder] = None, db_name: str = "aiogram_fsm", - states_collection_name: str = "states", - data_collection_name: str = "data", + 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 states_collection_name: name of the collection for storing FSM states. - :param data_collection_name: name of the collection for storing additional data. + :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._db_name = db_name - self._states_collection: AsyncIOMotorCollection = self._client[db_name][ - states_collection_name - ] - self._data_collection: AsyncIOMotorCollection = self._client[db_name][data_collection_name] + self._collection: AsyncIOMotorCollection = self._client[db_name][collection_name] self._key_builder = key_builder @classmethod @@ -71,50 +66,65 @@ class MongoStorage(BaseStorage): return str(value) async def set_state(self, key: StorageKey, state: StateType = None) -> None: - document_id = self._key_builder.build(key, "state") + document_id = self._key_builder.build(key) if state is None: - await self._states_collection.delete_one({"_id": document_id}) + 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._states_collection.update_one( - {"_id": document_id}, - {"$set": {"state": self.resolve_state(state)}}, + 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, "state") - document = await self._states_collection.find_one({"_id": document_id}) - if document is None or document["state"] is None: + document_id = self._key_builder.build(key) + document = await self._collection.find_one({"_id": document_id}) + if document is None or document.get("state") is None: return None return str(document["state"]) async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None: - document_id = self._key_builder.build(key, "data") + document_id = self._key_builder.build(key) if not data: - await self._data_collection.delete_one({"_id": document_id}) + 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._data_collection.update_one( - {"_id": document_id}, - {"$set": data}, + 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, "data") - document = await self._data_collection.find_one({"_id": document_id}, {"_id": 0}) - if not document: + 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) + 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, "data") - update_result = await self._data_collection.find_one_and_update( - {"_id": document_id}, - {"$set": data}, + 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._data_collection.delete_one({"_id": document_id}) - return cast(Dict[str, Any], update_result) + await self._collection.delete_one({"_id": document_id}) + return update_result.get("data", {}) diff --git a/tests/test_fsm/storage/test_key_builder.py b/tests/test_fsm/storage/test_key_builder.py index 7fc808d0..a459a564 100644 --- a/tests/test_fsm/storage/test_key_builder.py +++ b/tests/test_fsm/storage/test_key_builder.py @@ -1,3 +1,5 @@ +from typing import Literal, Optional + import pytest from aiogram.fsm.storage.base import DEFAULT_DESTINY, StorageKey @@ -13,22 +15,29 @@ FIELD = "data" class TestDefaultKeyBuilder: @pytest.mark.parametrize( - "with_bot_id,with_destiny,result", + "with_bot_id,with_destiny,field,result", [ - [False, False, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"], - [True, False, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}"], - [True, True, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"], - [False, True, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"], + [False, False, FIELD, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"], + [True, False, FIELD, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}"], + [True, True, FIELD, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"], + [False, True, FIELD, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"], + [False, False, None, f"{PREFIX}:{CHAT_ID}:{USER_ID}"], ], ) - async def test_generate_key(self, with_bot_id: bool, with_destiny: bool, result: str): + async def test_generate_key( + self, + with_bot_id: bool, + with_destiny: bool, + field: Optional[Literal["data", "state", "lock"]], + result: str, + ): key_builder = DefaultKeyBuilder( prefix=PREFIX, with_bot_id=with_bot_id, with_destiny=with_destiny, ) key = StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=BOT_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( diff --git a/tests/test_fsm/storage/test_mongo.py b/tests/test_fsm/storage/test_mongo.py index dfecd9cf..48b95719 100644 --- a/tests/test_fsm/storage/test_mongo.py +++ b/tests/test_fsm/storage/test_mongo.py @@ -4,20 +4,147 @@ 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=-42, user_id=42, bot_id=bot.id) + 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._data_collection.find_one({}) is None + 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._data_collection.find_one({}) is None + 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(