diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 6fe9188d3..ced26e58d 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -174,7 +174,7 @@ class AutomatedMaster(IMaster): current_depth = self._current_depth if self._current_depth is not None else 0 logger.info(f"Current depth is {current_depth}") - if should_propagate(self._control_channel.get_config(), self._current_depth): + if should_propagate(config.propagation.maximum_depth, self._current_depth): self._propagator.propagate(config.propagation, current_depth, self._stop) else: logger.info("Skipping propagation: maximum depth reached") diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index cf9283526..23a5a3ace 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -175,7 +175,9 @@ class InfectionMonkey: self._monkey_inbound_tunnel = self._control_client.create_control_tunnel( config.keep_tunnel_open_time ) - if self._monkey_inbound_tunnel and should_propagate(config, self._current_depth): + if self._monkey_inbound_tunnel and should_propagate( + config.propagation.maximum_depth, self._current_depth + ): self._inbound_tunnel_opened = True self._monkey_inbound_tunnel.start() diff --git a/monkey/infection_monkey/utils/propagation.py b/monkey/infection_monkey/utils/propagation.py index 9cfccab51..5036b5745 100644 --- a/monkey/infection_monkey/utils/propagation.py +++ b/monkey/infection_monkey/utils/propagation.py @@ -1,5 +1,2 @@ -from common.configuration import AgentConfiguration - - -def should_propagate(config: AgentConfiguration, current_depth: int) -> bool: - return config.propagation.maximum_depth > current_depth +def should_propagate(maximum_depth: int, current_depth: int) -> bool: + return maximum_depth > current_depth 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 15c15d0f0..a4531f92d 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py @@ -1,47 +1,22 @@ -import pytest - -from common.configuration import AgentConfiguration, AgentConfigurationSchema from infection_monkey.utils.propagation import should_propagate -@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(get_config): - max_depth = 2 +def test_should_propagate_current_less_than_max(): + maximum_depth = 2 current_depth = 1 - config = get_config(max_depth) - - assert should_propagate(config, current_depth) is True + assert should_propagate(maximum_depth, current_depth) is True -def test_should_propagate_current_greater_than_max(get_config): - max_depth = 2 +def test_should_propagate_current_greater_than_max(): + maximum_depth = 2 current_depth = 3 - config = get_config(max_depth) - - assert should_propagate(config, current_depth) is False + assert should_propagate(maximum_depth, current_depth) is False -def test_should_propagate_current_equal_to_max(get_config): - max_depth = 2 - current_depth = max_depth +def test_should_propagate_current_equal_to_max(): + maximum_depth = 2 + current_depth = maximum_depth - config = get_config(max_depth) - - assert should_propagate(config, current_depth) is False + assert should_propagate(maximum_depth, current_depth) is False