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

UT: Pass MagicMock instead of instance of ControlClient

This commit is contained in:
Ilija Lazoroski 2022-06-14 16:49:45 +02:00
parent e6e6587f46
commit 5ff617b811
2 changed files with 11 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py
View File

@ -2,7 +2,6 @@ from unittest.mock import MagicMock
import pytest import pytest
from infection_monkey.control import ControlClient
from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA
MONKEY_DIR_PATH = "/dir/to/monkey/" MONKEY_DIR_PATH = "/dir/to/monkey/"
@ -47,15 +46,8 @@ def fake_custom_pba_linux_options():
} }
@pytest.fixture def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux):
def control_client(): pba = CustomPBA(MagicMock(), MagicMock())
return ControlClient(CUSTOM_SERVER)
def test_command_linux_custom_file_and_cmd(
fake_custom_pba_linux_options, set_os_linux, control_client
):
pba = CustomPBA(MagicMock(), control_client)
pba._set_options(fake_custom_pba_linux_options) pba._set_options(fake_custom_pba_linux_options)
expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}" expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}"
assert pba.command == expected_command assert pba.command == expected_command
@ -72,11 +64,9 @@ def fake_custom_pba_windows_options():
} }
def test_command_windows_custom_file_and_cmd( def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows):
fake_custom_pba_windows_options, set_os_windows, control_client
):
pba = CustomPBA(MagicMock(), control_client) pba = CustomPBA(MagicMock(), MagicMock())
pba._set_options(fake_custom_pba_windows_options) pba._set_options(fake_custom_pba_windows_options)
expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}" expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}"
assert pba.command == expected_command assert pba.command == expected_command
@ -94,8 +84,8 @@ def fake_options_files_only():
@pytest.mark.parametrize("os", [set_os_linux, set_os_windows]) @pytest.mark.parametrize("os", [set_os_linux, set_os_windows])
def test_files_only(fake_options_files_only, os, control_client): def test_files_only(fake_options_files_only, os):
pba = CustomPBA(MagicMock(), control_client) pba = CustomPBA(MagicMock(), MagicMock())
pba._set_options(fake_options_files_only) pba._set_options(fake_options_files_only)
assert pba.command == "" assert pba.command == ""
@ -110,15 +100,15 @@ def fake_options_commands_only():
} }
def test_commands_only(fake_options_commands_only, set_os_linux, control_client): def test_commands_only(fake_options_commands_only, set_os_linux):
pba = CustomPBA(MagicMock(), control_client) pba = CustomPBA(MagicMock(), MagicMock())
pba._set_options(fake_options_commands_only) pba._set_options(fake_options_commands_only)
assert pba.command == CUSTOM_LINUX_CMD assert pba.command == CUSTOM_LINUX_CMD
assert pba.filename == "" assert pba.filename == ""
def test_commands_only_windows(fake_options_commands_only, set_os_windows, control_client): def test_commands_only_windows(fake_options_commands_only, set_os_windows):
pba = CustomPBA(MagicMock(), control_client) pba = CustomPBA(MagicMock(), MagicMock())
pba._set_options(fake_options_commands_only) pba._set_options(fake_options_commands_only)
assert pba.command == CUSTOM_WINDOWS_CMD assert pba.command == CUSTOM_WINDOWS_CMD
assert pba.filename == "" assert pba.filename == ""

3
monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py
View File

@ -13,7 +13,6 @@ def spy_send_telemetry(monkeypatch):
_spy_send_telemetry.telem_category = None _spy_send_telemetry.telem_category = None
_spy_send_telemetry.data = None _spy_send_telemetry.data = None
control_client = ControlClient("localhost:5000") ControlClient.control_client_object = MagicMock()
ControlClient.control_client_object = control_client
ControlClient.control_client_object.send_telemetry = MagicMock(side_effect=_spy_send_telemetry) ControlClient.control_client_object.send_telemetry = MagicMock(side_effect=_spy_send_telemetry)
return _spy_send_telemetry return _spy_send_telemetry