Tests: Reduced code duplication in database initializer tests

This commit is contained in:
Mike Salvatore 2021-09-30 13:13:26 -04:00
parent f0c25b4b5e
commit fef6350871
1 changed files with 38 additions and 29 deletions

View File

@ -1,4 +1,3 @@
from pathlib import Path
from unittest.mock import MagicMock from unittest.mock import MagicMock
import mongomock import mongomock
@ -8,49 +7,59 @@ from monkey_island.cc.setup.mongo.database_initializer import reset_database
@pytest.fixture @pytest.fixture
def fake_mongo(monkeypatch): def patch_attack_mitigations_path(monkeypatch, data_for_tests_dir):
mongo = mongomock.MongoClient() def inner(file_name):
monkeypatch.setattr("monkey_island.cc.setup.mongo.database_initializer.mongo", mongo) path = data_for_tests_dir / "mongo_mitigations" / file_name
monkeypatch.setattr("monkey_island.cc.services.database.mongo", mongo) monkeypatch.setattr(
return mongo "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 @pytest.fixture
def fake_config(monkeypatch): def mock_mongo_client(monkeypatch):
monkeypatch.setattr("monkey_island.cc.services.config.ConfigService.init_config", lambda: None) mongo = mongomock.MongoClient()
monkeypatch.setattr("monkey_island.cc.services.attack.attack_config.AttackConfig.reset_config", lambda: None) mongo.db.validate_collection = MagicMock(return_value=True)
monkeypatch.setattr("monkey_island.cc.services.database.jsonify", 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): def test_store_mitigations_on_mongo(patch_attack_mitigations_path, mock_mongo_client):
monkeypatch.setattr( patch_attack_mitigations_path("attack_mitigations.json")
"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)
reset_database() 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( def test_store_mitigations_on_mongo__invalid_mitigation(patch_attack_mitigations_path):
monkeypatch, data_for_tests_dir, fake_mongo, fake_config patch_attack_mitigations_path("invalid_mitigation")
):
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)
with pytest.raises(Exception): with pytest.raises(Exception):
reset_database() reset_database()
def test_get_all_mitigations(monkeypatch, fake_mongo, fake_config): def test_get_all_mitigations(mock_mongo_client):
fake_mongo.db.validate_collection = MagicMock(return_value=True)
reset_database() reset_database()
mitigations = list(fake_mongo.db.attack_mitigations.find({})) mitigations = list(mock_mongo_client.db.attack_mitigations.find({}))
assert len(mitigations) >= 266 assert len(mitigations) >= 266