diff --git a/.gitignore b/.gitignore index d1dbf6d6..1ce11ad0 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,4 @@ reports dev/ .venv/ +.conda/ diff --git a/aiogram/fsm/storage/mongo.py b/aiogram/fsm/storage/mongo.py index b4b1eeaa..420eea5e 100644 --- a/aiogram/fsm/storage/mongo.py +++ b/aiogram/fsm/storage/mongo.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional, cast +from typing import Any, Dict, Optional, cast, overload from motor.motor_asyncio import AsyncIOMotorClient @@ -115,6 +115,23 @@ class MongoStorage(BaseStorage): return {} return cast(Dict[str, Any], document["data"]) + @overload + async def get_value(self, storage_key: StorageKey, dict_key: str) -> Optional[Any]: ... + + @overload + async def get_value(self, storage_key: StorageKey, dict_key: str, default: Any) -> Any: ... + + async def get_value( + self, storage_key: StorageKey, dict_key: str, default: Optional[Any] = None + ) -> Optional[Any]: + document_id = self._key_builder.build(storage_key) + projection = {"_id": 0, f"data.{dict_key}": 1} + document = await self._collection.find_one({"_id": document_id}, projection=projection) + if not document or "data" not in document: + return default + nested_data = document.get("data", {}) + return nested_data.get(dict_key, default) + 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()} diff --git a/tests/test_fsm/storage/test_mongo.py b/tests/test_fsm/storage/test_mongo.py index 95a476d6..23c8acb5 100644 --- a/tests/test_fsm/storage/test_mongo.py +++ b/tests/test_fsm/storage/test_mongo.py @@ -149,6 +149,88 @@ class TestStateAndDataDoNotAffectEachOther: } +class TestGetValue: + async def test_get_existing_value( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_data( + storage_key, {"key": "value", "number": 42, "list": [1, 2, 3]} + ) + + assert await mongo_storage.get_value(storage_key, "key") == "value" + assert await mongo_storage.get_value(storage_key, "number") == 42 + assert await mongo_storage.get_value(storage_key, "list") == [1, 2, 3] + + async def test_get_non_existing_value( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_data(storage_key, {"key": "value"}) + + assert await mongo_storage.get_value(storage_key, "non_existing_key") is None + assert ( + await mongo_storage.get_value(storage_key, "non_existing_key", default="default") + == "default" + ) + + async def test_get_value_from_non_existing_document( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + assert await mongo_storage._collection.find_one({}) is None + + assert await mongo_storage.get_value(storage_key, "any_key") is None + assert ( + await mongo_storage.get_value(storage_key, "any_key", default="default") == "default" + ) + + async def test_get_value_with_document_containing_only_state( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + ): + await mongo_storage.set_state(storage_key, "test") + + document = await mongo_storage._collection.find_one({}) + assert document is not None + assert "data" not in document + + assert await mongo_storage.get_value(storage_key, "any_key") is None + assert ( + await mongo_storage.get_value(storage_key, "any_key", default="default") == "default" + ) + + async def test_get_value_uses_projection( + self, + mongo_storage: MongoStorage, + storage_key: StorageKey, + monkeypatch, + ): + await mongo_storage.set_data( + storage_key, {"key1": "value1", "key2": "value2", "key3": {"nested": "data"}} + ) + + original_find_one = mongo_storage._collection.find_one + calls = [] + + async def mock_find_one(*args, **kwargs): + calls.append(kwargs) + return await original_find_one(*args, **kwargs) + + monkeypatch.setattr(mongo_storage._collection, "find_one", mock_find_one) + + value = await mongo_storage.get_value(storage_key, "key2") + + assert len(calls) == 1 + assert "projection" in calls[0] + assert calls[0]["projection"] == {"_id": 0, "data.key2": 1} + assert value == "value2" + + @pytest.mark.parametrize( "value,result", [