From 5be0a3d6f9a73c7129b87a7616877b15ff308b40 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 14 Mar 2022 12:38:06 -0400 Subject: [PATCH] UT: Use a mock IAgentRepository instead of monkeypatching open() --- .../infection_monkey/exploit/test_powershell.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py index 6d0aa6df3..d3d516353 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -1,3 +1,4 @@ +from io import BytesIO from unittest.mock import MagicMock import pytest @@ -21,6 +22,10 @@ class AuthenticationErrorForTests(Exception): pass +mock_agent_repository = MagicMock() +mock_agent_repository.get_agent_binary.return_value = BytesIO(b"BINARY_EXECUTABLE") + + @pytest.fixture def powershell_arguments(): options = { @@ -37,7 +42,7 @@ def powershell_arguments(): "options": options, "current_depth": 2, "telemetry_messenger": MagicMock(), - "agent_repository": MagicMock(), + "agent_repository": mock_agent_repository, } return arguments @@ -48,7 +53,6 @@ def powershell_exploiter(monkeypatch): monkeypatch.setattr(powershell, "AuthenticationError", AuthenticationErrorForTests) monkeypatch.setattr(powershell, "is_windows_os", lambda: True) - monkeypatch.setattr("builtins.open", lambda _, __: MagicMock()) return pe @@ -160,6 +164,7 @@ def test_successful_propagation(monkeypatch, powershell_exploiter, powershell_ar assert exploit_result.exploitation_success assert exploit_result.propagation_success + assert not exploit_result.error_message def test_login_attempts_correctly_reported(monkeypatch, powershell_exploiter, powershell_arguments):