Added test for the db_from_uri parameter in MongoStorage; enabled generic test TestStorage for mongo

This commit is contained in:
Grigory Statsenko 2022-07-13 00:21:35 +03:00
parent 0ab6b15e5d
commit 7e233737c8
4 changed files with 72 additions and 1 deletions

View file

@ -21,6 +21,7 @@ STATE = 'aiogram_state'
DATA = 'aiogram_data'
BUCKET = 'aiogram_bucket'
COLLECTIONS = (STATE, DATA, BUCKET)
DEFAULT_DB_NAME = 'aiogram_fsm'
class MongoStorage(BaseStorage):
@ -42,7 +43,7 @@ class MongoStorage(BaseStorage):
await dp.storage.wait_closed()
"""
def __init__(self, host='localhost', port=27017, db_name='aiogram_fsm', uri=None,
def __init__(self, host='localhost', port=27017, db_name=DEFAULT_DB_NAME, uri=None,
username=None, password=None, index=True, db_from_uri=False, **kwargs):
self._host = host
self._port = port

View file

@ -17,3 +17,4 @@ rethinkdb>=2.4.1
coverage==4.5.3
motor>=2.2.0
pytest-lazy-fixture==0.6.*
yarl

View file

@ -3,6 +3,8 @@ import asyncio
import aioredis
import pytest
from _pytest.config import UsageError
from motor.motor_asyncio import AsyncIOMotorClient
from yarl import URL
from aiogram import Bot
from . import TOKEN
@ -19,6 +21,14 @@ def pytest_addoption(parser):
default=None,
help="run tests which require redis connection",
)
parser.addoption(
"--mongo",
default=None,
help=(
"run tests which require mongo connection, e.g.: "
"mongodb://test:qwerty@127.0.0.1:27017/test_db?authSource=admin"
),
)
def pytest_configure(config):
@ -26,6 +36,10 @@ def pytest_configure(config):
"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):
@ -79,6 +93,27 @@ def redis_options(request):
raise UsageError("Unsupported aioredis version")
@pytest.fixture(scope='session')
def mongo_options(request):
mongo_uri = request.config.getoption("--mongo")
if mongo_uri is None:
pytest.skip("need --mongo option with mongo URI to run")
return
mongo_client = AsyncIOMotorClient(mongo_uri)
db_name = mongo_client.get_default_database().name
# MongoDB URI-s are pretty much standard URLs so we can use yarl.URL() on them
url = URL(mongo_uri)
host, port = url.host, url.port
options = {
'host': host,
'port': port,
'db_name': db_name,
'uri': mongo_uri,
}
return options
@pytest.fixture(name='bot')
async def bot_fixture():
"""Bot fixture."""

View file

@ -3,6 +3,7 @@ import pytest
from pytest_lazyfixture import lazy_fixture
from aiogram.contrib.fsm_storage.memory import MemoryStorage
from aiogram.contrib.fsm_storage.redis import RedisStorage, RedisStorage2
from aiogram.contrib.fsm_storage.mongo import MongoStorage, DEFAULT_DB_NAME
@pytest.fixture()
@ -39,11 +40,18 @@ async def memory_store():
yield MemoryStorage()
@pytest.fixture()
@pytest.mark.mongo
async def mongo_store(mongo_options):
yield MongoStorage(uri=mongo_options['uri'], db_from_uri=True)
@pytest.mark.parametrize(
"store", [
lazy_fixture('redis_store'),
lazy_fixture('redis_store2'),
lazy_fixture('memory_store'),
lazy_fixture('mongo_store'),
]
)
class TestStorage:
@ -82,3 +90,29 @@ class TestRedisStorage2:
assert await store.get_data(chat='1234') == {
'foo': 'bar'} # new pool was opened at this point
assert id(store._redis) != pool_id
@pytest.mark.asyncio
@pytest.mark.mongo
async def test_mongo_dbname_from_uri(mongo_options):
"""
Check that it is possible to use the db_name provided in the URI
without specifying it as a separate parameter.
"""
mongo_uri = mongo_options['uri']
default_db_name = DEFAULT_DB_NAME
uri_db_name = mongo_options['db_name']
assert uri_db_name in mongo_uri
assert default_db_name not in mongo_uri
# Make two instances - one with the parameter, and the other - without
mongo_store_1 = MongoStorage(uri=mongo_options['uri'])
mongo_store_2 = MongoStorage(uri=mongo_options['uri'], db_from_uri=True)
db_1 = await mongo_store_1.get_db()
db_2 = await mongo_store_2.get_db()
assert db_1.name == default_db_name
assert db_2.name == uri_db_name