python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2026-04-08 16:37:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Storing FSM state and data together in MongoDB-storage

Browse source
This commit is contained in:
Rishat Fayzullin 2024-03-13 21:36:57 +03:00
parent 7c18c87b95
commit ce48b825e8
4 changed files with 199 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

23
tests/test_fsm/storage/test_key_builder.py
View file

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

133
tests/test_fsm/storage/test_mongo.py
View file

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