mirror of
https://github.com/aiogram/aiogram.git
synced 2025-12-15 03:08:51 +00:00
Added MongoStorage for FSM (#1434)
* 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
GitHub
parent
25c76b7d74
commit
1ef7655fd7
15 changed files with 522 additions and 129 deletions
|
|
@ -2,6 +2,8 @@ from pathlib import Path
|
|||
|
||||
import pytest
|
||||
from _pytest.config import UsageError
|
||||
from pymongo.errors import InvalidURI, PyMongoError
|
||||
from pymongo.uri_parser import parse_uri as parse_mongo_url
|
||||
from redis.asyncio.connection import parse_url as parse_redis_url
|
||||
|
||||
from aiogram import Dispatcher
|
||||
|
|
@ -10,6 +12,7 @@ from aiogram.fsm.storage.memory import (
|
|||
MemoryStorage,
|
||||
SimpleEventIsolation,
|
||||
)
|
||||
from aiogram.fsm.storage.mongo import MongoStorage
|
||||
from aiogram.fsm.storage.redis import RedisEventIsolation, RedisStorage
|
||||
from tests.mocked_bot import MockedBot
|
||||
|
||||
|
|
@ -18,24 +21,27 @@ DATA_DIR = Path(__file__).parent / "data"
|
|||
|
||||
def pytest_addoption(parser):
|
||||
parser.addoption("--redis", default=None, help="run tests which require redis connection")
|
||||
parser.addoption("--mongo", default=None, help="run tests which require mongo connection")
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line("markers", "redis: marked tests require redis connection to run")
|
||||
config.addinivalue_line("markers", "mongo: marked tests require mongo connection to run")
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
redis_uri = config.getoption("--redis")
|
||||
if redis_uri is None:
|
||||
skip_redis = pytest.mark.skip(reason="need --redis option with redis URI to run")
|
||||
for item in items:
|
||||
if "redis" in item.keywords:
|
||||
item.add_marker(skip_redis)
|
||||
return
|
||||
try:
|
||||
parse_redis_url(redis_uri)
|
||||
except ValueError as e:
|
||||
raise UsageError(f"Invalid redis URI {redis_uri!r}: {e}")
|
||||
for db, parse_uri in [("redis", parse_redis_url), ("mongo", parse_mongo_url)]:
|
||||
uri = config.getoption(f"--{db}")
|
||||
if uri is None:
|
||||
skip = pytest.mark.skip(reason=f"need --{db} option with {db} URI to run")
|
||||
for item in items:
|
||||
if db in item.keywords:
|
||||
item.add_marker(skip)
|
||||
else:
|
||||
try:
|
||||
parse_uri(uri)
|
||||
except (ValueError, InvalidURI) as e:
|
||||
raise UsageError(f"Invalid {db} URI {uri!r}: {e}")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
|
|
@ -62,6 +68,29 @@ async def redis_storage(redis_server):
|
|||
await storage.close()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mongo_server(request):
|
||||
mongo_uri = request.config.getoption("--mongo")
|
||||
return mongo_uri
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
@pytest.mark.mongo
|
||||
async def mongo_storage(mongo_server):
|
||||
if not mongo_server:
|
||||
pytest.skip("MongoDB is not available here")
|
||||
storage = MongoStorage.from_url(mongo_server)
|
||||
try:
|
||||
await storage._client.server_info()
|
||||
except PyMongoError as e:
|
||||
pytest.skip(str(e))
|
||||
else:
|
||||
yield storage
|
||||
await storage._client.drop_database(storage._database)
|
||||
finally:
|
||||
await storage.close()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
async def memory_storage():
|
||||
storage = MemoryStorage()
|
||||
|
|
|
|||
|
|
@ -5,3 +5,11 @@ services:
|
|||
image: redis:6-alpine
|
||||
ports:
|
||||
- "${REDIS_PORT-6379}:6379"
|
||||
|
||||
mongo:
|
||||
image: mongo:7.0.6
|
||||
environment:
|
||||
MONGO_INITDB_ROOT_USERNAME: mongo
|
||||
MONGO_INITDB_ROOT_PASSWORD: mongo
|
||||
ports:
|
||||
- "${MONGODB_PORT-27017}:27017"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
164
tests/test_fsm/storage/test_mongo.py
Normal 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
|
||||
|
|
@ -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",
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue