From 1ba10d70595c01104a7c730ecb6dd26e9cc3b687 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 9 Sep 2021 10:33:29 +0200 Subject: [PATCH] UT: Fix powershell copy_file tests --- .../infection_monkey/exploit/test_powershell.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 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 fa24ee7ed..7c2a04710 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -118,25 +118,23 @@ def authenticate(mock_client): ) def test_successful_copy(monkeypatch, powershell_exploiter, dropper_target_path, arch): mock_client = MagicMock() - mock_client.get_host_architecture = lambda: arch - mock_client.copy_file = MagicMock(return_value=True) + mock_client.return_value.get_host_architecture = lambda: arch + mock_client.return_value.copy_file = MagicMock(return_value=True) - mock_powershell_client = MagicMock(side_effect=authenticate(mock_client)) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + monkeypatch.setattr(powershell, "PowerShellClient", mock_client) success = powershell_exploiter.exploit_host() - assert dropper_target_path in mock_client.copy_file.call_args[0][1] + assert dropper_target_path in mock_client.return_value.copy_file.call_args[0][1] assert success def test_failed_copy(monkeypatch, powershell_exploiter): mock_client = MagicMock() - mock_client.get_host_architecture = lambda: WIN_ARCH_32 - mock_client.copy_file = MagicMock(return_value=False) + mock_client.return_value.get_host_architecture = lambda: WIN_ARCH_32 + mock_client.return_value.copy_file = MagicMock(return_value=False) - mock_powershell_client = MagicMock(side_effect=authenticate(mock_client)) - monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) + monkeypatch.setattr(powershell, "PowerShellClient", mock_client) success = powershell_exploiter.exploit_host() assert not success