python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2025-12-16 20:23:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added MongoStorage for FSM (#1434)

Browse source 
* Mongo storage included to storages test

* Added few additional checks in storages test

* Added MongoStorage for FSM

* Added changes description

* Fixed error message syntax

Co-authored-by: Alex Root Junior <jroot.junior@gmail.com>

* Resolved mypy check error

* IF/ELSE statement simplified

* Fix ruff linter error: RET505 Unnecessary `elif` after `return` statement

* Fix ruff linter error: E501 Line too long (100 > 99)

* Added mongo storage testing in CI

* Refactoring while review

* Refactoring while review

* Storing FSM state and data together in MongoDB-storage

* Fix CI - MongoDB container action is only supported on Linux

* Refactoring while review

* Enable Macos in pypy-tests section of CI

* Refactoring while review

* Makefile updated

* redis and mongo storages tests do not run in pypy-tests job of CI

* Fix docstring of DefaultKeyBuilder

---------

Co-authored-by: Alex Root Junior <jroot.junior@gmail.com>
This commit is contained in:
Rishat-F 2024-05-07 22:42:31 +03:00 • committed by GitHub
parent 25c76b7d74
commit 1ef7655fd7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 522 additions and 129 deletions
Download patch file Download diff file Expand all files Collapse all files

10
tests/test_fsm/storage/test_isolation.py
View file

@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch
import pytest
from aiogram.fsm.storage.base import BaseEventIsolation, StorageKey
from aiogram.fsm.storage.redis import RedisEventIsolation
from aiogram.fsm.storage.redis import RedisEventIsolation, RedisStorage
from tests.mocked_bot import MockedBot
@ -32,6 +32,14 @@ class TestIsolations:
class TestRedisEventIsolation:
def test_create_isolation(self):
fake_redis = object()
storage = RedisStorage(redis=fake_redis)
isolation = storage.create_isolation()
assert isinstance(isolation, RedisEventIsolation)
assert isolation.redis is fake_redis
assert isolation.key_builder is storage.key_builder
def test_init_without_key_builder(self):
redis = AsyncMock()
isolation = RedisEventIsolation(redis=redis)

45
tests/test_fsm/storage/test_redis.py → tests/test_fsm/storage/test_key_builder.py
View file

@ -1,11 +1,8 @@
from typing import Literal, Optional
import pytest
from aiogram.fsm.storage.base import DEFAULT_DESTINY, StorageKey
from aiogram.fsm.storage.redis import (
DefaultKeyBuilder,
RedisEventIsolation,
RedisStorage,
)
from aiogram.fsm.storage.base import DEFAULT_DESTINY, DefaultKeyBuilder, StorageKey
PREFIX = "test"
BOT_ID = 42
@ -16,9 +13,9 @@ BUSINESS_CONNECTION_ID = "4"
FIELD = "data"
class TestRedisDefaultKeyBuilder:
class TestDefaultKeyBuilder:
@pytest.mark.parametrize(
"key_builder,result",
"key_builder,field,result",
[
[
DefaultKeyBuilder(
@ -27,40 +24,52 @@ class TestRedisDefaultKeyBuilder:
with_destiny=True,
with_business_connection_id=True,
),
FIELD,
f"{PREFIX}:{BOT_ID}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}",
],
[
DefaultKeyBuilder(prefix=PREFIX, with_bot_id=True, with_destiny=True),
f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}",
None,
f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}",
],
[
DefaultKeyBuilder(
prefix=PREFIX, with_bot_id=True, with_business_connection_id=True
),
FIELD,
f"{PREFIX}:{BOT_ID}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{FIELD}",
],
[
DefaultKeyBuilder(prefix=PREFIX, with_bot_id=True),
f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}",
None,
f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}",
],
[
DefaultKeyBuilder(
prefix=PREFIX, with_destiny=True, with_business_connection_id=True
),
FIELD,
f"{PREFIX}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}",
],
[
DefaultKeyBuilder(prefix=PREFIX, with_destiny=True),
f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}",
None,
f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}",
],
[
DefaultKeyBuilder(prefix=PREFIX, with_business_connection_id=True),
FIELD,
f"{PREFIX}:{BUSINESS_CONNECTION_ID}:{CHAT_ID}:{USER_ID}:{FIELD}",
],
[DefaultKeyBuilder(prefix=PREFIX), f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"],
[DefaultKeyBuilder(prefix=PREFIX), None, f"{PREFIX}:{CHAT_ID}:{USER_ID}"],
],
)
async def test_generate_key(self, key_builder: DefaultKeyBuilder, result: str):
async def test_generate_key(
self,
key_builder: DefaultKeyBuilder,
field: Optional[Literal["data", "state", "lock"]],
result: str,
):
key = StorageKey(
chat_id=CHAT_ID,
user_id=USER_ID,
@ -68,7 +77,7 @@ class TestRedisDefaultKeyBuilder:
business_connection_id=BUSINESS_CONNECTION_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(
@ -95,11 +104,3 @@ class TestRedisDefaultKeyBuilder:
destiny=DEFAULT_DESTINY,
)
assert key_builder.build(key, FIELD) == f"{PREFIX}:{CHAT_ID}:{THREAD_ID}:{USER_ID}:{FIELD}"
def test_create_isolation(self):
fake_redis = object()
storage = RedisStorage(redis=fake_redis)
isolation = storage.create_isolation()
assert isinstance(isolation, RedisEventIsolation)
assert isolation.redis is fake_redis
assert isolation.key_builder is storage.key_builder

164
tests/test_fsm/storage/test_mongo.py Normal file
View file

@ -0,0 +1,164 @@
import pytest
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=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._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._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(
"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, mongo_storage: MongoStorage):
assert mongo_storage.resolve_state(value) == result

16
tests/test_fsm/storage/test_storages.py
View file

@ -11,7 +11,11 @@ def create_storage_key(bot: MockedBot):
@pytest.mark.parametrize(
"storage",
[pytest.lazy_fixture("redis_storage"), pytest.lazy_fixture("memory_storage")],
[
pytest.lazy_fixture("redis_storage"),
pytest.lazy_fixture("mongo_storage"),
pytest.lazy_fixture("memory_storage"),
],
)
class TestStorages:
async def test_set_state(self, bot: MockedBot, storage: BaseStorage, storage_key: StorageKey):
@ -35,6 +39,8 @@ class TestStorages:
):
assert await storage.get_data(key=storage_key) == {}
assert await storage.update_data(key=storage_key, data={"foo": "bar"}) == {"foo": "bar"}
assert await storage.update_data(key=storage_key, data={}) == {"foo": "bar"}
assert await storage.get_data(key=storage_key) == {"foo": "bar"}
assert await storage.update_data(key=storage_key, data={"baz": "spam"}) == {
"foo": "bar",
"baz": "spam",
@ -43,3 +49,11 @@ class TestStorages:
"foo": "bar",
"baz": "spam",
}
assert await storage.update_data(key=storage_key, data={"baz": "test"}) == {
"foo": "bar",
"baz": "test",
}
assert await storage.get_data(key=storage_key) == {
"foo": "bar",
"baz": "test",
}