mirror of
https://github.com/aiogram/aiogram.git
synced 2026-04-08 16:37:47 +00:00
feat(storage): add get_value method to MongoStorage
Implement get_value method for MongoStorage that retrieves individual values from storage data using MongoDB projections. Add comprehensive test coverage for the new method.
This commit is contained in:
parent
25e9127db9
commit
49e9b34b3f
3 changed files with 101 additions and 1 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -22,3 +22,4 @@ reports
|
|||
|
||||
dev/
|
||||
.venv/
|
||||
.conda/
|
||||
|
|
|
|||
|
|
@ -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()}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue