p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

UT: Use a mock IAgentRepository instead of monkeypatching open()

This commit is contained in:
Mike Salvatore 2022-03-14 12:38:06 -04:00 committed by Ilija Lazoroski
parent 020dbbf2fe
commit 5be0a3d6f9
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py
View File

@ -1,3 +1,4 @@
from io import BytesIO
from unittest.mock import MagicMock from unittest.mock import MagicMock
import pytest import pytest
@ -21,6 +22,10 @@ class AuthenticationErrorForTests(Exception):
pass pass
mock_agent_repository = MagicMock()
mock_agent_repository.get_agent_binary.return_value = BytesIO(b"BINARY_EXECUTABLE")
@pytest.fixture @pytest.fixture
def powershell_arguments(): def powershell_arguments():
options = { options = {
@ -37,7 +42,7 @@ def powershell_arguments():
"options": options, "options": options,
"current_depth": 2, "current_depth": 2,
"telemetry_messenger": MagicMock(), "telemetry_messenger": MagicMock(),
"agent_repository": MagicMock(), "agent_repository": mock_agent_repository,
} }
return arguments return arguments
@ -48,7 +53,6 @@ def powershell_exploiter(monkeypatch):
monkeypatch.setattr(powershell, "AuthenticationError", AuthenticationErrorForTests) monkeypatch.setattr(powershell, "AuthenticationError", AuthenticationErrorForTests)
monkeypatch.setattr(powershell, "is_windows_os", lambda: True) monkeypatch.setattr(powershell, "is_windows_os", lambda: True)
monkeypatch.setattr("builtins.open", lambda _, __: MagicMock())
return pe return pe
@ -160,6 +164,7 @@ def test_successful_propagation(monkeypatch, powershell_exploiter, powershell_ar
assert exploit_result.exploitation_success assert exploit_result.exploitation_success
assert exploit_result.propagation_success assert exploit_result.propagation_success
assert not exploit_result.error_message
def test_login_attempts_correctly_reported(monkeypatch, powershell_exploiter, powershell_arguments): def test_login_attempts_correctly_reported(monkeypatch, powershell_exploiter, powershell_arguments):