UT: Replace test data with SocketAddress wherever required

This commit is contained in:
Shreya Malviya 2022-09-26 18:36:08 +05:30
parent 0ae653fb72
commit 3b192a869e
5 changed files with 15 additions and 9 deletions

View File

@ -11,6 +11,7 @@ from common.agent_event_serializers import (
) )
from common.agent_events import AbstractAgentEvent from common.agent_events import AbstractAgentEvent
from common.agent_registration_data import AgentRegistrationData from common.agent_registration_data import AgentRegistrationData
from common.types import SocketAddress
from infection_monkey.island_api_client import ( from infection_monkey.island_api_client import (
HTTPIslandAPIClient, HTTPIslandAPIClient,
IslandAPIConnectionError, IslandAPIConnectionError,
@ -20,7 +21,7 @@ from infection_monkey.island_api_client import (
IslandAPITimeoutError, IslandAPITimeoutError,
) )
SERVER = "1.1.1.1:9999" SERVER = SocketAddress(ip="1.1.1.1", port="9999")
PBA_FILE = "dummy.pba" PBA_FILE = "dummy.pba"
WINDOWS = "windows" WINDOWS = "windows"
AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673") AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673")

View File

@ -7,6 +7,7 @@ from uuid import UUID
import pytest import pytest
from common import AgentRegistrationData from common import AgentRegistrationData
from common.types import SocketAddress
from monkey_island.cc.island_event_handlers import handle_agent_registration from monkey_island.cc.island_event_handlers import handle_agent_registration
from monkey_island.cc.models import Agent, CommunicationType, Machine from monkey_island.cc.models import Agent, CommunicationType, Machine
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
@ -31,7 +32,7 @@ AGENT_REGISTRATION_DATA = AgentRegistrationData(
machine_hardware_id=MACHINE.hardware_id, machine_hardware_id=MACHINE.hardware_id,
start_time=0, start_time=0,
parent_id=None, parent_id=None,
cc_server="192.168.1.1:5000", cc_server=SocketAddress(ip="192.168.1.1", port="5000"),
network_interfaces=[IPv4Interface("192.168.1.2/24")], network_interfaces=[IPv4Interface("192.168.1.2/24")],
) )
@ -111,7 +112,7 @@ def test_existing_machine_updated__find_by_ip(handler, machine_repository):
machine_hardware_id=5, machine_hardware_id=5,
start_time=0, start_time=0,
parent_id=None, parent_id=None,
cc_server="192.168.1.1:5000", cc_server=SocketAddress(ip="192.168.1.1", port="5000"),
network_interfaces=[ network_interfaces=[
IPv4Interface("192.168.1.2/24"), IPv4Interface("192.168.1.2/24"),
IPv4Interface("192.168.1.4/24"), IPv4Interface("192.168.1.4/24"),
@ -215,7 +216,7 @@ def test_machine_interfaces_updated(handler, machine_repository):
machine_hardware_id=MACHINE.hardware_id, machine_hardware_id=MACHINE.hardware_id,
start_time=0, start_time=0,
parent_id=None, parent_id=None,
cc_server="192.168.1.1:5000", cc_server=SocketAddress(ip="192.168.1.1", port="5000"),
network_interfaces=[ network_interfaces=[
IPv4Interface("192.168.1.2/24"), IPv4Interface("192.168.1.2/24"),
IPv4Interface("192.168.1.3/16"), IPv4Interface("192.168.1.3/16"),

View File

@ -6,16 +6,18 @@ from uuid import UUID
import pytest import pytest
from common import AgentRegistrationData from common import AgentRegistrationData
from common.types import SocketAddress
AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10") AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10")
PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484") PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484")
SOCKET_ADDRESS = SocketAddress(ip="192.168.1.1", port="5000")
AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = { AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = {
"id": AGENT_ID, "id": AGENT_ID,
"machine_hardware_id": 2, "machine_hardware_id": 2,
"start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc), "start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc),
"parent_id": PARENT_ID, "parent_id": PARENT_ID,
"cc_server": "192.168.1.1:5000", "cc_server": SOCKET_ADDRESS,
"network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")], "network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")],
} }
@ -24,7 +26,7 @@ AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT = {
"machine_hardware_id": 2, "machine_hardware_id": 2,
"start_time": "2022-08-18T18:46:48+00:00", "start_time": "2022-08-18T18:46:48+00:00",
"parent_id": str(PARENT_ID), "parent_id": str(PARENT_ID),
"cc_server": "192.168.1.1:5000", "cc_server": SOCKET_ADDRESS.dict(simplify=True),
"network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"], "network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"],
} }
@ -85,7 +87,7 @@ def test_construct_invalid_field__value_error(key, value):
("machine_hardware_id", 99), ("machine_hardware_id", 99),
("start_time", 0), ("start_time", 0),
("parent_id", AGENT_ID), ("parent_id", AGENT_ID),
("cc_server", "10.0.0.1:4999"), ("cc_server", SOCKET_ADDRESS),
("network_interfaces", ["10.0.0.1/24"]), ("network_interfaces", ["10.0.0.1/24"]),
], ],
) )

View File

@ -6,6 +6,7 @@ from uuid import uuid4
import mongomock import mongomock
import pytest import pytest
from common.types import SocketAddress
from monkey_island.cc.models import Agent from monkey_island.cc.models import Agent
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
IAgentRepository, IAgentRepository,
@ -54,6 +55,7 @@ AGENTS = (
*RUNNING_AGENTS, *RUNNING_AGENTS,
*STOPPED_AGENTS, *STOPPED_AGENTS,
) )
CC_SERVER = SocketAddress(ip="127.0.0.1", port="1984")
@pytest.fixture @pytest.fixture
@ -116,7 +118,7 @@ def test_upsert_agent__insert_empty_repository(empty_agent_repository):
def test_upsert_agent__update(agent_repository): def test_upsert_agent__update(agent_repository):
agents = deepcopy(AGENTS) agents = deepcopy(AGENTS)
agents[0].stop_time = datetime.now() agents[0].stop_time = datetime.now()
agents[0].cc_server = "127.0.0.1:1984" agents[0].cc_server = CC_SERVER
agent_repository.upsert_agent(agents[0]) agent_repository.upsert_agent(agents[0])

View File

@ -16,7 +16,7 @@ AGENT_REGISTRATION_DICT = {
"machine_hardware_id": 1, "machine_hardware_id": 1,
"start_time": 0, "start_time": 0,
"parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"), "parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"),
"cc_server": "10.0.0.1:5000", "cc_server": {"ip": "10.0.0.1", "port": "5000"},
"network_interfaces": ["10.1.1.2/24"], "network_interfaces": ["10.1.1.2/24"],
} }