From afeca66d92df422ee581a02ce67b0fc1fa78b355 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 23 Jun 2022 16:27:05 -0400 Subject: [PATCH] UT: Use AgentConfiguration in test_propagation.py --- .../utils/test_propagation.py | 25 +++++++++++++++---- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py b/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py index 37f7194a6..15c15d0f0 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py @@ -1,11 +1,26 @@ +import pytest + +from common.configuration import AgentConfiguration, AgentConfigurationSchema from infection_monkey.utils.propagation import should_propagate -def get_config(max_depth): - return {"config": {"depth": max_depth}} +@pytest.fixture +def get_config(default_agent_configuration): + def _inner(max_depth): + # AgentConfiguration is a frozen dataclass, so we need to deserialize and reserialize to + # modify it. The benefit is that it's impossible to construct an invalid object. The + # downside is the extra steps required to change an object. Maybe we can come up with a + # better all-around solution. It depends how often we need to mutate these objects (probably + # only for tests). + agent_dict = AgentConfigurationSchema().dump(default_agent_configuration) + agent_dict["propagation"]["maximum_depth"] = max_depth + + return AgentConfiguration.from_dict(agent_dict) + + return _inner -def test_should_propagate_current_less_than_max(): +def test_should_propagate_current_less_than_max(get_config): max_depth = 2 current_depth = 1 @@ -14,7 +29,7 @@ def test_should_propagate_current_less_than_max(): assert should_propagate(config, current_depth) is True -def test_should_propagate_current_greater_than_max(): +def test_should_propagate_current_greater_than_max(get_config): max_depth = 2 current_depth = 3 @@ -23,7 +38,7 @@ def test_should_propagate_current_greater_than_max(): assert should_propagate(config, current_depth) is False -def test_should_propagate_current_equal_to_max(): +def test_should_propagate_current_equal_to_max(get_config): max_depth = 2 current_depth = max_depth