From 65d43575d17603681f9861539d6cc580d908c7eb Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 28 Sep 2022 16:04:49 +0000 Subject: [PATCH] UT: Make tests pass --- .../handle_ping_scan_event.py | 10 +-- .../test_handle_ping_scan_event.py | 84 ++++++++++++++++++- 2 files changed, 86 insertions(+), 8 deletions(-) diff --git a/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py b/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py index 52fc34ddc..bae20b607 100644 --- a/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py +++ b/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py @@ -39,13 +39,13 @@ class handle_ping_scan_event: logger.exception("Unable to process ping scan data") def _get_destination_machine(self, event: PingScanEvent) -> Machine: - dest_machines = self._machine_repository.get_machines_by_ip(event.target) - if not dest_machines: + try: + dest_machines = self._machine_repository.get_machines_by_ip(event.target) + return dest_machines[0] + except UnknownRecordError: machine = Machine(id=self._machine_repository.get_new_id()) - dest_machines = [machine] self._machine_repository.upsert_machine(machine) - - return dest_machines[0] + return machine def _get_source_machine(self, event: PingScanEvent) -> Machine: agent = self._agent_repository.get_agent_by_id(event.source) diff --git a/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py b/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py index 171266ae3..ba1a7ff60 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py +++ b/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py @@ -14,6 +14,8 @@ from monkey_island.cc.repository import ( IAgentRepository, IMachineRepository, INodeRepository, + RetrievalError, + StorageError, UnknownRecordError, ) @@ -68,10 +70,51 @@ def handler(agent_repository, machine_repository, node_repository) -> handle_pin machines = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} +machines_by_id = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} +machines_by_ip = { + IPv4Address("10.10.10.99"): [PINGER_MACHINE], + IPv4Address("10.10.10.1"): [TARGET_MACHINE], +} def machine_from_id(id: int): - return machines[id] + return machines_by_id[id] + + +def machines_from_ip(ip: IPv4Address): + return machines_by_ip[ip] + + +class error_machine_by_id: + """Raise an error if the machine with the called ID matches the stored ID""" + + def __init__(self, id: int, error): + self.id = id + self.error = error + + def __call__(self, id: int): + if id == self.id: + raise self.error + else: + return machine_from_id(id) + + +class error_machine_by_ip: + """Raise an error if the machine with the called IP matches the stored ID""" + + def __init__(self, id: int, error): + self.id = id + self.error = error + + def __call__(self, ip: IPv4Address): + print(f"IP is: {ip}") + machines = machines_from_ip(ip) + if machines[0].id == self.id: + print(f"Raise error: {self.error}") + raise self.error + else: + print(f"Return machine: {machines}") + return machines def test_handle_ping_scan_event__upserts_machine( @@ -79,6 +122,7 @@ def test_handle_ping_scan_event__upserts_machine( machine_repository: IMachineRepository, ): machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) handler(EVENT) expected_machine = TARGET_MACHINE.copy() @@ -91,7 +135,8 @@ 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 _: []) + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) handler(EVENT) expected_machine = TARGET_MACHINE.copy() @@ -105,7 +150,8 @@ def test_handle_ping_scan_event__upserts_node( machine_repository: IMachineRepository, node_repository: INodeRepository, ): - machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE) + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE]) handler(EVENT) node_repository.upsert_communication.assert_called_with( @@ -143,8 +189,40 @@ def test_handle_ping_scan_event__upserts_machine_if_not_existed( handler: handle_ping_scan_event, machine_repository: IMachineRepository ): machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) handler(EVENT) expected_machine = Machine(id=SEED_ID, operating_system=OperatingSystem.LINUX) machine_repository.upsert_machine.assert_called_with(expected_machine) + + +@pytest.mark.parametrize("id", [PINGER_MACHINE.id, TARGET_MACHINE.id]) +def test_handle_scan_data__node_not_upserted_if_machine_retrievalerror( + handler: handle_ping_scan_event, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + id, +): + machine_repository.get_machine_by_id = MagicMock( + side_effect=error_machine_by_id(id, RetrievalError) + ) + machine_repository.get_machines_by_ip = MagicMock( + side_effect=error_machine_by_ip(id, RetrievalError) + ) + + handler(EVENT) + + assert not node_repository.upsert_communication.called + + +def test_handle_scan_data__node_not_upserted_if_machine_storageerror( + handler: handle_ping_scan_event, + machine_repository: IMachineRepository, + node_repository: INodeRepository, +): + machine_repository.upsert_machine = MagicMock(side_effect=StorageError) + + handler(EVENT) + + assert not node_repository.upsert_communication.called