diff --git a/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_database_initializer.py b/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_database_initializer.py index ed20c5ea0..d3ca3fbcc 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_database_initializer.py +++ b/monkey/tests/unit_tests/monkey_island/cc/setup/mongo/test_database_initializer.py @@ -1,4 +1,3 @@ -from pathlib import Path from unittest.mock import MagicMock import mongomock @@ -8,49 +7,59 @@ from monkey_island.cc.setup.mongo.database_initializer import reset_database @pytest.fixture -def fake_mongo(monkeypatch): - mongo = mongomock.MongoClient() - monkeypatch.setattr("monkey_island.cc.setup.mongo.database_initializer.mongo", mongo) - monkeypatch.setattr("monkey_island.cc.services.database.mongo", mongo) - return mongo +def patch_attack_mitigations_path(monkeypatch, data_for_tests_dir): + def inner(file_name): + path = data_for_tests_dir / "mongo_mitigations" / file_name + monkeypatch.setattr( + "monkey_island.cc.setup.mongo.database_initializer.ATTACK_MITIGATION_PATH", path + ) + + return inner + + +@pytest.fixture(scope="module", autouse=True) +def patch_dependencies(monkeypatch_session): + monkeypatch_session.setattr( + "monkey_island.cc.services.config.ConfigService.init_config", lambda: None + ) + monkeypatch_session.setattr( + "monkey_island.cc.services.attack.attack_config.AttackConfig.reset_config", lambda: None + ) + monkeypatch_session.setattr( + "monkey_island.cc.services.database.jsonify", MagicMock(return_value=True) + ) @pytest.fixture -def fake_config(monkeypatch): - monkeypatch.setattr("monkey_island.cc.services.config.ConfigService.init_config", lambda: None) - monkeypatch.setattr("monkey_island.cc.services.attack.attack_config.AttackConfig.reset_config", lambda: None) - monkeypatch.setattr("monkey_island.cc.services.database.jsonify", MagicMock(return_value=True)) +def mock_mongo_client(monkeypatch): + mongo = mongomock.MongoClient() + mongo.db.validate_collection = MagicMock(return_value=True) + + monkeypatch.setattr("monkey_island.cc.setup.mongo.database_initializer.mongo", mongo) + monkeypatch.setattr("monkey_island.cc.services.database.mongo", mongo) + + return mongo -def test_store_mitigations_on_mongo(monkeypatch, data_for_tests_dir, fake_mongo, fake_config): - monkeypatch.setattr( - "monkey_island.cc.setup.mongo.database_initializer.ATTACK_MITIGATION_PATH", - Path(data_for_tests_dir) / "mongo_mitigations" / "attack_mitigations.json", - ) - fake_mongo.db.validate_collection = MagicMock(return_value=True) +def test_store_mitigations_on_mongo(patch_attack_mitigations_path, mock_mongo_client): + patch_attack_mitigations_path("attack_mitigations.json") + reset_database() - assert len(list(fake_mongo.db.attack_mitigations.find({}))) == 3 + assert len(list(mock_mongo_client.db.attack_mitigations.find({}))) == 3 -def test_store_mitigations_on_mongo__invalid_mitigation( - monkeypatch, data_for_tests_dir, fake_mongo, fake_config -): - monkeypatch.setattr( - "monkey_island.cc.setup.mongo.database_initializer.ATTACK_MITIGATION_PATH", - Path(data_for_tests_dir) / "mongo_mitigations" / "invalid_mitigation", - ) - fake_mongo.db.validate_collection = MagicMock(return_value=True) +def test_store_mitigations_on_mongo__invalid_mitigation(patch_attack_mitigations_path): + patch_attack_mitigations_path("invalid_mitigation") + with pytest.raises(Exception): reset_database() -def test_get_all_mitigations(monkeypatch, fake_mongo, fake_config): - fake_mongo.db.validate_collection = MagicMock(return_value=True) - +def test_get_all_mitigations(mock_mongo_client): reset_database() - mitigations = list(fake_mongo.db.attack_mitigations.find({})) + mitigations = list(mock_mongo_client.db.attack_mitigations.find({})) assert len(mitigations) >= 266