forked from p34709852/monkey
Agent: Decouple should_propagate() and AgentConfiguration
This commit is contained in:
parent
afeca66d92
commit
ad0f6946bd
|
@ -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")
|
||||
|
|
|
@ -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()
|
||||
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue