From ad0f6946bdad1ffbb8038bc299acd14d673da195 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 23 Jun 2022 16:31:28 -0400
Subject: [PATCH] Agent: Decouple should_propagate() and AgentConfiguration

---
 .../master/automated_master.py                |  2 +-
 monkey/infection_monkey/monkey.py             |  4 +-
 monkey/infection_monkey/utils/propagation.py  |  7 +--
 .../utils/test_propagation.py                 | 45 +++++--------------
 4 files changed, 16 insertions(+), 42 deletions(-)

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