diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py
index e7da336eb..2618f8d93 100644
--- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py
+++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py
@@ -1,3 +1,5 @@
+from unittest.mock import MagicMock
+
 import pytest
 
 from infection_monkey.post_breach.actions.users_custom_pba import UsersPBA
@@ -43,7 +45,7 @@ def mock_UsersPBA_linux_custom_file_and_cmd(set_os_linux, fake_monkey_dir_path,
         "infection_monkey.config.WormConfiguration.PBA_linux_filename",
         CUSTOM_LINUX_FILENAME,
     )
-    return UsersPBA()
+    return UsersPBA(MagicMock())
 
 
 def test_command_linux_custom_file_and_cmd(
@@ -63,7 +65,7 @@ def mock_UsersPBA_windows_custom_file_and_cmd(set_os_windows, fake_monkey_dir_pa
         "infection_monkey.config.WormConfiguration.PBA_windows_filename",
         CUSTOM_WINDOWS_FILENAME,
     )
-    return UsersPBA()
+    return UsersPBA(MagicMock())
 
 
 def test_command_windows_custom_file_and_cmd(
@@ -80,7 +82,7 @@ def mock_UsersPBA_linux_custom_file(set_os_linux, fake_monkey_dir_path, monkeypa
         "infection_monkey.config.WormConfiguration.PBA_linux_filename",
         CUSTOM_LINUX_FILENAME,
     )
-    return UsersPBA()
+    return UsersPBA(MagicMock())
 
 
 def test_command_linux_custom_file(mock_UsersPBA_linux_custom_file):
@@ -95,7 +97,7 @@ def mock_UsersPBA_windows_custom_file(set_os_windows, fake_monkey_dir_path, monk
         "infection_monkey.config.WormConfiguration.PBA_windows_filename",
         CUSTOM_WINDOWS_FILENAME,
     )
-    return UsersPBA()
+    return UsersPBA(MagicMock())
 
 
 def test_command_windows_custom_file(mock_UsersPBA_windows_custom_file):
@@ -110,7 +112,7 @@ def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypat
         CUSTOM_LINUX_CMD,
     )
     monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_linux_filename", None)
-    return UsersPBA()
+    return UsersPBA(MagicMock())
 
 
 def test_command_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd):
@@ -125,7 +127,7 @@ def mock_UsersPBA_windows_custom_cmd(set_os_windows, fake_monkey_dir_path, monke
         CUSTOM_WINDOWS_CMD,
     )
     monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_windows_filename", None)
-    return UsersPBA()
+    return UsersPBA(MagicMock())
 
 
 def test_command_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd):
diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_post_breach_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_post_breach_telem.py
index d71a82e2a..fff070ad1 100644
--- a/monkey/tests/unit_tests/infection_monkey/telemetry/test_post_breach_telem.py
+++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_post_breach_telem.py
@@ -2,6 +2,7 @@ import json
 
 import pytest
 
+from infection_monkey.i_puppet import PostBreachData
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 
 HOSTNAME = "hostname"
@@ -22,7 +23,7 @@ class StubSomePBA:
 def post_breach_telem_test_instance(monkeypatch):
     monkeypatch.setattr(PostBreachTelem, "_get_hostname_and_ip", lambda: (HOSTNAME, IP))
     monkeypatch.setattr(PostBreachTelem, "_get_os", lambda: OS)
-    return PostBreachTelem(PBA_NAME, PBA_COMMAND, RESULT)
+    return PostBreachTelem(PostBreachData(PBA_NAME, PBA_COMMAND, RESULT))
 
 
 def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry):