feat: Remove/Keep only specified DBs in mongodb_factory

newsfragments/812.feature.rst

@@ -0,0 +1,3 @@
+Improve mongodb_factory to keep or delete provided databases.
+It's useful to run xdist tests in parallel, and avoid deleting databases created by other test workers.
+Now you can keep preseed databases which won't be deleted after each test.

pytest_mongo/config.py

@@ -16,6 +16,8 @@ class MongoConfig:
     port_search_count: int
     params: str
     tz_aware: bool
+    remove_dbs: list[str] | None
+    keep_dbs: list[str] | None
 
 
 def get_config(request: FixtureRequest) -> MongoConfig:
@@ -34,5 +36,7 @@ def get_mongo_option(option: str) -> Any:
         port_search_count=int(get_mongo_option("port_search_count")),
         params=get_mongo_option("params"),
         tz_aware=get_mongo_option("tz_aware"),
+        remove_dbs=get_mongo_option("remove_dbs"),
+        keep_dbs=get_mongo_option("keep_dbs"),
     )
     return cfg

pytest_mongo/factories/client.py

@@ -10,15 +10,26 @@
 
 
 def mongodb(
-    process_fixture_name: str, tz_aware: bool | None = None
+    process_fixture_name: str,
+    tz_aware: bool | None = None,
+    remove_dbs: list[str] | None = None,
+    keep_dbs: list[str] | None = None,
 ) -> Callable[[FixtureRequest], Iterator[MongoClient]]:
     """Mongo database factory.
 
     :param str process_fixture_name: name of the process fixture
     :param bool tz_aware: whether the client to be timezone aware or not
+    :param list remove_dbs: list of database names to be cleaned
+    :param list keep_dbs: list of database names to be kept
     :rtype: func
     :returns: function which makes a connection to mongo
     """
+    if (
+        remove_dbs is not None
+        and keep_dbs is not None
+        and set(remove_dbs).intersection(set(keep_dbs))
+    ):
+        raise ValueError("remove_dbs and keep_dbs cannot have overlapping database names")
 
     @pytest.fixture
     def mongodb_factory(request: FixtureRequest) -> Iterator[MongoClient]:
@@ -36,6 +47,18 @@ def mongodb_factory(request: FixtureRequest) -> Iterator[MongoClient]:
         elif config.tz_aware is not None and isinstance(config.tz_aware, bool):
             mongo_tz_aware = config.tz_aware
 
+        mongo_remove_dbs = None
+        if remove_dbs is not None:
+            mongo_remove_dbs = remove_dbs
+        elif config.remove_dbs is not None and isinstance(config.remove_dbs, list):
+            mongo_remove_dbs = config.remove_dbs
+
+        mongo_keep_dbs = None
+        if keep_dbs is not None:
+            mongo_keep_dbs = keep_dbs
+        elif config.keep_dbs is not None and isinstance(config.keep_dbs, list):
+            mongo_keep_dbs = config.keep_dbs
+
         mongo_host = mongodb_process.host
         mongo_port = mongodb_process.port
 
@@ -44,6 +67,12 @@ def mongodb_factory(request: FixtureRequest) -> Iterator[MongoClient]:
         yield mongo_conn
 
         for db_name in mongo_conn.list_database_names():
+            if mongo_keep_dbs and db_name in mongo_keep_dbs:
+                continue
+
+            if mongo_remove_dbs and db_name not in mongo_remove_dbs:
+                continue
+
             database = mongo_conn[db_name]
             for collection_name in database.list_collection_names():
                 collection = database[collection_name]

pytest_mongo/plugin.py

@@ -31,6 +31,14 @@
     "Have mongo client timezone aware (ini: mongo_tz_aware; "
     "use --mongo-tz-aware/--no-mongo-tz-aware to override)"
 )
+_help_remove_dbs = (
+    "List of MongoDB databases to clean in the fixture. Otherwise, all databases "
+    "are cleaned excluding system.*. collections"
+)
+_help_keep_dbs = (
+    "List of MongoDB databases to keep in the fixture. Otherwise, all databases "
+    "are cleaned excluding system.*. collections"
+)
 
 
 def pytest_addoption(parser: Parser) -> None:
@@ -57,6 +65,20 @@ def pytest_addoption(parser: Parser) -> None:
         default=False,
     )
 
+    parser.addini(
+        name="mongo_remove_dbs",
+        type="args",
+        help=_help_remove_dbs,
+        default=[],
+    )
+
+    parser.addini(
+        name="mongo_keep_dbs",
+        type="args",
+        help=_help_keep_dbs,
+        default=[],
+    )
+
     parser.addoption(
         "--mongo-exec",
         action="store",
@@ -95,6 +117,22 @@ def pytest_addoption(parser: Parser) -> None:
         dest="mongo_tz_aware",
         help=_help_tz_aware,
     )
+    parser.addoption(
+        "--mongo-remove-dbs",
+        help=_help_remove_dbs,
+        dest="mongo_remove_dbs",
+        type=str,
+        nargs="*",
+        action="extend",
+    )
+    parser.addoption(
+        "--mongo-keep-dbs",
+        help=_help_keep_dbs,
+        dest="mongo_keep_dbs",
+        type=str,
+        nargs="*",
+        action="extend",
+    )
 
 
 mongo_proc = factories.mongo_proc()

tests/conftest.py

@@ -13,4 +13,8 @@
 
 mongo_proc_rand = process.mongo_proc(port=None, params=mongo_params)
 mongodb_rand = client.mongodb("mongo_proc_rand")
+
+mongo_proc4 = process.mongo_proc(port=27072, params=mongo_params)
+mongodb4 = client.mongodb("mongo_proc4", remove_dbs=["test_db"])
+mongodb5 = client.mongodb("mongo_proc4", keep_dbs=["test_db2"])
 # pylint:enable=invalid-name

tests/test_mongo_cleanup.py

@@ -0,0 +1,26 @@
+"""Test MongoDB cleanup functionality."""
+
+from pymongo import MongoClient
+
+
+def test_clean_specified_databases(mongodb4: MongoClient) -> None:
+    """Test if only specified databases are cleaned."""
+    test_db = mongodb4["test_db"]
+    test_db.test.insert_one({"test": "test"})
+    test_db2 = mongodb4["test_db2"]
+    test_db2.test.insert_one({"test": "test"})
+
+    assert "test_db" in mongodb4.list_database_names()
+    assert "test_db2" in mongodb4.list_database_names()
+
+
+def test_clean_specified_databases_again(mongodb5: MongoClient) -> None:
+    """Test if only specified databases are cleaned."""
+    assert "test_db" not in mongodb5.list_database_names()
+    assert "test_db2" in mongodb5.list_database_names()
+
+
+def test_keep_specified_databases(mongodb5: MongoClient) -> None:
+    """Test if only specified databases are kept."""
+    assert "test_db" not in mongodb5.list_database_names()
+    assert "test_db2" in mongodb5.list_database_names()