UT: Use AgentConfiguration in test_propagation.py

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