diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py
index 92aff10fc..48f8b6bf5 100644
--- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py
+++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock
 import pytest
 
 from monkey.infection_monkey.network.relay import (
-    RELAY_CONTROL_MESSAGE,
+    RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST,
     RelayConnectionHandler,
     RelayUserHandler,
     TCPPipeSpawner,
@@ -27,7 +27,7 @@ def relay_user_handler():
 @pytest.fixture
 def close_socket():
     sock = MagicMock(spec=socket.socket)
-    sock.recv.return_value = RELAY_CONTROL_MESSAGE
+    sock.recv.return_value = RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST
     sock.getpeername.return_value = (USER_ADDRESS, 12345)
     return sock