UT: Fix control client and custom pba tests

monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py

@@ -2,6 +2,7 @@ from unittest.mock import MagicMock
 
 import pytest
 
+from infection_monkey.control import ControlClient
 from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA
 
 MONKEY_DIR_PATH = "/dir/to/monkey/"
@@ -48,8 +49,15 @@ def fake_custom_pba_linux_options():
     }
 
 
-def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux):
-    pba = CustomPBA(MagicMock())
+@pytest.fixture
+def control_client():
+    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)
     expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}"
     assert pba.command == expected_command
@@ -68,9 +76,11 @@ def fake_custom_pba_windows_options():
     }
 
 
-def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows):
+def test_command_windows_custom_file_and_cmd(
+    fake_custom_pba_windows_options, set_os_windows, control_client
+):
 
-    pba = CustomPBA(MagicMock())
+    pba = CustomPBA(MagicMock(), control_client)
     pba._set_options(fake_custom_pba_windows_options)
     expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}"
     assert pba.command == expected_command
@@ -90,8 +100,8 @@ def fake_options_files_only():
 
 
 @pytest.mark.parametrize("os", [set_os_linux, set_os_windows])
-def test_files_only(fake_options_files_only, os):
-    pba = CustomPBA(MagicMock())
+def test_files_only(fake_options_files_only, os, control_client):
+    pba = CustomPBA(MagicMock(), control_client)
     pba._set_options(fake_options_files_only)
     assert pba.command == ""
 
@@ -108,15 +118,15 @@ def fake_options_commands_only():
     }
 
 
-def test_commands_only(fake_options_commands_only, set_os_linux):
-    pba = CustomPBA(MagicMock())
+def test_commands_only(fake_options_commands_only, set_os_linux, control_client):
+    pba = CustomPBA(MagicMock(), control_client)
     pba._set_options(fake_options_commands_only)
     assert pba.command == CUSTOM_LINUX_CMD
     assert pba.filename == ""
 
 
-def test_commands_only_windows(fake_options_commands_only, set_os_windows):
-    pba = CustomPBA(MagicMock())
+def test_commands_only_windows(fake_options_commands_only, set_os_windows, control_client):
+    pba = CustomPBA(MagicMock(), control_client)
     pba._set_options(fake_options_commands_only)
     assert pba.command == CUSTOM_WINDOWS_CMD
     assert pba.filename == ""

monkey/tests/unit_tests/infection_monkey/test_control.py

@@ -8,8 +8,8 @@ from monkey.infection_monkey.control import ControlClient
     [(True, "http://8.8.8.8:45455"), (False, "8.8.8.8:45455")],
 )
 def test_control_set_proxies(monkeypatch, is_windows_os, expected_proxy_string):
-    monkeypatch.setattr("monkey.infection_monkey.control.is_windows_os", lambda: is_windows_os)
-    control_client = ControlClient()
+    monkeypatch.setattr("infection_monkey.cc.control.is_windows_os", lambda: is_windows_os)
+    control_client = ControlClient("8.8.8.8:5000")
 
     control_client.set_proxies(("8.8.8.8", "45455"))