diff --git a/newsfragments/812.feature.rst b/newsfragments/812.feature.rst new file mode 100644 index 0000000..d6d16b5 --- /dev/null +++ b/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. diff --git a/pytest_mongo/config.py b/pytest_mongo/config.py index be52a6f..6434b45 100644 --- a/pytest_mongo/config.py +++ b/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 diff --git a/pytest_mongo/factories/client.py b/pytest_mongo/factories/client.py index 971cfea..d1b1943 100644 --- a/pytest_mongo/factories/client.py +++ b/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] diff --git a/pytest_mongo/plugin.py b/pytest_mongo/plugin.py index d0a8b1a..8ff5698 100644 --- a/pytest_mongo/plugin.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py index 96179c7..4180eaf 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_mongo_cleanup.py b/tests/test_mongo_cleanup.py new file mode 100644 index 0000000..e195cc8 --- /dev/null +++ b/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()