diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py b/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py new file mode 100644 index 000000000..37f7194a6 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_propagation.py @@ -0,0 +1,32 @@ +from infection_monkey.utils.propagation import should_propagate + + +def get_config(max_depth): + return {"config": {"depth": max_depth}} + + +def test_should_propagate_current_less_than_max(): + max_depth = 2 + current_depth = 1 + + config = get_config(max_depth) + + assert should_propagate(config, current_depth) is True + + +def test_should_propagate_current_greater_than_max(): + max_depth = 2 + current_depth = 3 + + config = get_config(max_depth) + + assert should_propagate(config, current_depth) is False + + +def test_should_propagate_current_equal_to_max(): + max_depth = 2 + current_depth = max_depth + + config = get_config(max_depth) + + assert should_propagate(config, current_depth) is False