forked from p15670423/monkey
UT: Use RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST in place of RELAY_CONTROL_MESSAGE
This commit is contained in:
parent
aa1c31efb2
commit
451d2d0694
|
@ -5,7 +5,7 @@ from unittest.mock import MagicMock
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from monkey.infection_monkey.network.relay import (
|
from monkey.infection_monkey.network.relay import (
|
||||||
RELAY_CONTROL_MESSAGE,
|
RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST,
|
||||||
RelayConnectionHandler,
|
RelayConnectionHandler,
|
||||||
RelayUserHandler,
|
RelayUserHandler,
|
||||||
TCPPipeSpawner,
|
TCPPipeSpawner,
|
||||||
|
@ -27,7 +27,7 @@ def relay_user_handler():
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def close_socket():
|
def close_socket():
|
||||||
sock = MagicMock(spec=socket.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)
|
sock.getpeername.return_value = (USER_ADDRESS, 12345)
|
||||||
return sock
|
return sock
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue