From 7e233737c87715c66c26488e439a1b84c99826e8 Mon Sep 17 00:00:00 2001
From: Grigory Statsenko <grisha100@gmail.com>
Date: Wed, 13 Jul 2022 00:21:35 +0300
Subject: [PATCH] Added test for the db_from_uri parameter in MongoStorage;
 enabled generic test TestStorage for mongo

---
 aiogram/contrib/fsm_storage/mongo.py      |  3 +-
 dev_requirements.txt                      |  1 +
 tests/conftest.py                         | 35 +++++++++++++++++++++++
 tests/contrib/fsm_storage/test_storage.py | 34 ++++++++++++++++++++++
 4 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/aiogram/contrib/fsm_storage/mongo.py b/aiogram/contrib/fsm_storage/mongo.py
index 0f873b1f..a4481d53 100644
--- a/aiogram/contrib/fsm_storage/mongo.py
+++ b/aiogram/contrib/fsm_storage/mongo.py
@@ -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
diff --git a/dev_requirements.txt b/dev_requirements.txt
index 26e410aa..41400f70 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -17,3 +17,4 @@ rethinkdb>=2.4.1
 coverage==4.5.3
 motor>=2.2.0
 pytest-lazy-fixture==0.6.*
+yarl
diff --git a/tests/conftest.py b/tests/conftest.py
index 455f8b0d..d97eb8ef 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -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."""
diff --git a/tests/contrib/fsm_storage/test_storage.py b/tests/contrib/fsm_storage/test_storage.py
index 2668cdab..5f281e9f 100644
--- a/tests/contrib/fsm_storage/test_storage.py
+++ b/tests/contrib/fsm_storage/test_storage.py
@@ -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