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

UT: Fix and rename unit tests in test_handle_ping_scan_event.py

This commit is contained in:
vakarisz 2022-09-28 18:04:10 +03:00 committed by Kekoa Kaaikala
parent b6d9f88dee
commit 168a5845fd
1 changed files with 28 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_scan_data.py → monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py
View File

@ -14,8 +14,6 @@ from monkey_island.cc.repository import (
IAgentRepository, IAgentRepository,
IMachineRepository, IMachineRepository,
INodeRepository, INodeRepository,
RetrievalError,
StorageError,
UnknownRecordError, UnknownRecordError,
) )
@ -24,12 +22,12 @@ AGENT_ID = UUID("1d8ce743-a0f4-45c5-96af-91106529d3e2")
MACHINE_ID = 11 MACHINE_ID = 11
CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000") CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000")
AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER) AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER)
MACHINE = Machine( PINGER_MACHINE = Machine(
id=MACHINE_ID, id=MACHINE_ID,
hardware_id=5, hardware_id=5,
network_interfaces=[IPv4Interface("10.10.10.99/24")], network_interfaces=[IPv4Interface("10.10.10.99/24")],
) )
STORED_MACHINE = Machine( TARGET_MACHINE = Machine(
id=33, id=33,
hardware_id=9, hardware_id=9,
network_interfaces=[IPv4Interface("10.10.10.1/24")], network_interfaces=[IPv4Interface("10.10.10.1/24")],
@ -69,54 +67,67 @@ def handler(agent_repository, machine_repository, node_repository) -> handle_pin
return handle_ping_scan_event(agent_repository, machine_repository, node_repository) return handle_ping_scan_event(agent_repository, machine_repository, node_repository)
machines = {MACHINE_ID: MACHINE, STORED_MACHINE.id: STORED_MACHINE} machines = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE}
def machine_from_id(id: int): def machine_from_id(id: int):
return machines[id] return machines[id]
def test_handle_scan_data__upserts_machine( def test_handle_ping_scan_event__upserts_machine(
handler: handle_ping_scan_event, handler: handle_ping_scan_event,
machine_repository: IMachineRepository, machine_repository: IMachineRepository,
): ):
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
handler(EVENT) handler(EVENT)
expected_machine = STORED_MACHINE.copy() expected_machine = TARGET_MACHINE.copy()
expected_machine.operating_system = OperatingSystem.LINUX expected_machine.operating_system = OperatingSystem.LINUX
assert machine_repository.upsert_machine.called_with(expected_machine) machine_repository.upsert_machine.assert_called_with(expected_machine)
def test_handle_scan_data__upserts_node( def test_handle_ping_scan_event__machine_already_exists(
handler: handle_ping_scan_event,
machine_repository: IMachineRepository,
):
machine_repository.get_machine_by_id = MagicMock(side_effect=lambda _: [])
handler(EVENT)
expected_machine = TARGET_MACHINE.copy()
expected_machine.operating_system = OperatingSystem.LINUX
machine_repository.upsert_machine.assert_called_with(expected_machine)
def test_handle_ping_scan_event__upserts_node(
handler: handle_ping_scan_event, handler: handle_ping_scan_event,
machine_repository: IMachineRepository, machine_repository: IMachineRepository,
node_repository: INodeRepository, node_repository: INodeRepository,
): ):
machine_repository.get_machine_by_id = MagicMock(return_value=STORED_MACHINE) machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE)
handler(EVENT) handler(EVENT)
assert node_repository.upsert_communication.called_with( node_repository.upsert_communication.assert_called_with(
MACHINE.id, STORED_MACHINE.id, CommunicationType.SCANNED PINGER_MACHINE.id, TARGET_MACHINE.id, CommunicationType.SCANNED
) )
def test_handle_scan_data__node_not_upserted_if_no_matching_agent( def test_handle_ping_scan_event__node_not_upserted_if_no_matching_agent(
handler: handle_ping_scan_event, handler: handle_ping_scan_event,
agent_repository: IAgentRepository, agent_repository: IAgentRepository,
machine_repository: IMachineRepository, machine_repository: IMachineRepository,
node_repository: INodeRepository, node_repository: INodeRepository,
): ):
agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError) agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError)
machine_repository.get_machine_by_id = MagicMock(return_value=STORED_MACHINE) machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE)
handler(EVENT) handler(EVENT)
assert not node_repository.upsert_communication.called assert not node_repository.upsert_communication.called
def test_handle_scan_data__node_not_upserted_if_no_matching_machine( def test_handle_ping_scan_event__node_not_upserted_if_no_matching_machine(
handler: handle_ping_scan_event, handler: handle_ping_scan_event,
machine_repository: IMachineRepository, machine_repository: IMachineRepository,
node_repository: INodeRepository, node_repository: INodeRepository,
@ -128,7 +139,7 @@ def test_handle_scan_data__node_not_upserted_if_no_matching_machine(
assert not node_repository.upsert_communication.called assert not node_repository.upsert_communication.called
def test_handle_scan_data__upserts_machine_if_not_existed( def test_handle_ping_scan_event__upserts_machine_if_not_existed(
handler: handle_ping_scan_event, machine_repository: IMachineRepository handler: handle_ping_scan_event, machine_repository: IMachineRepository
): ):
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
@ -136,4 +147,4 @@ def test_handle_scan_data__upserts_machine_if_not_existed(
expected_machine = Machine(id=SEED_ID, operating_system=OperatingSystem.LINUX) expected_machine = Machine(id=SEED_ID, operating_system=OperatingSystem.LINUX)
assert machine_repository.upsert_machine.called_with(expected_machine) machine_repository.upsert_machine.assert_called_with(expected_machine)