diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index db782744..432356b8 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -1,8 +1,8 @@ -from typing import Any, Dict +from typing import Any, Dict, Union from primaite import getLogger from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent -from primaite.simulator.network.hardware.base import Link, Node +from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort _LOGGER = getLogger(__name__) @@ -72,16 +72,42 @@ class NetworkContainer(SimComponent): del self.nodes[node.uuid] del node.parent # misleading? - def connect_nodes(self, node1: Node, node2: Node) -> None: - """TODO.""" - # I think we should not be forcing users to add and remove individual links. - # Clearly if a link exists between two nodes in the network, then the link is also part of the network. - # I'm just not sure how we ought to handle link creation as it requires an unoccupied interface on the node. - raise NotImplementedError + def connect_nodes(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None: + """Connect two nodes on the network by creating a link between an NIC/SwitchPort of each one. - def disconnect_nodes(self, node1: Node, node2: Node) -> None: - """TODO.""" - raise NotImplementedError + :param endpoint_a: The endpoint to which to connect the link on the first node + :type endpoint_a: Union[NIC, SwitchPort] + :param endpoint_b: The endpoint to which to connct the link on the second node + :type endpoint_b: Union[NIC, SwitchPort] + :raises RuntimeError: _description_ + """ + node_a = endpoint_a.parent + node_b = endpoint_b.parent + msg = "" + if node_a not in self: + msg = f"Cannot create a link to {endpoint_a} because the node is not in the network." + if node_b not in self: + msg = f"Cannot create a link to {endpoint_b} because the node is not in the network." + if node_a is node_b: + msg = f"Cannot link {endpoint_a} to {endpoint_b} because they belong to the same node." + if msg: + _LOGGER.error(msg) + raise RuntimeError(msg) + + link = Link(endpoint_a=endpoint_a, endpoint_b=endpoint_b, **kwargs) + self.links[link.uuid] = link + link.parent = self + + def remove_link(self, link: Link) -> None: + """Disconnect a link from the network. + + :param link: The link to be removed + :type link: Link + """ + link.endpoint_a.disconnect_link() + link.endpoint_b.disconnect_link() + del self.links[link.uuid] + del link.parent def __contains__(self, item: Any) -> bool: if isinstance(item, Node): diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 28e7693a..5b49f008 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -918,6 +918,7 @@ class Node(SimComponent): if nic.uuid not in self.nics: self.nics[nic.uuid] = nic nic.connected_node = self + nic.parent = self self.sys_log.info(f"Connected NIC {nic}") if self.operating_state == NodeOperatingState.ON: nic.enable() @@ -938,6 +939,7 @@ class Node(SimComponent): nic = self.nics.get(nic) if nic or nic.uuid in self.nics: self.nics.pop(nic.uuid) + del nic.parent nic.disable() self.sys_log.info(f"Disconnected NIC {nic}") else: diff --git a/tests/integration_tests/network/test_network_creation.py b/tests/integration_tests/network/test_network_creation.py index 482c188d..0ee827be 100644 --- a/tests/integration_tests/network/test_network_creation.py +++ b/tests/integration_tests/network/test_network_creation.py @@ -1,7 +1,7 @@ import pytest from primaite.simulator.network.container import NetworkContainer -from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.hardware.base import NIC, Node def test_adding_removing_nodes(): @@ -36,3 +36,47 @@ def test_removing_nonexistent_node(): net.remove_node(n1) assert n1.parent is None assert n1 not in net + + +def test_connecting_nodes(): + """Check that two nodes on the network can be connected.""" + net = NetworkContainer() + n1 = Node(hostname="computer") + n1_nic = NIC(ip_address="120.30.0.1", gateway="192.168.0.1", subnet_mask="255.255.255.0") + n1.connect_nic(n1_nic) + n2 = Node(hostname="server") + n2_nic = NIC(ip_address="120.30.0.2", gateway="192.168.0.1", subnet_mask="255.255.255.0") + n2.connect_nic(n2_nic) + + net.add_node(n1) + net.add_node(n2) + + net.connect_nodes(n1.nics[n1_nic.uuid], n2.nics[n2_nic.uuid], bandwidth=30) + + assert len(net.links) == 1 + link = list(net.links.values())[0] + assert link in net + assert link.parent is net + + +def test_connecting_node_to_itself(): + net = NetworkContainer() + node = Node(hostname="computer") + nic1 = NIC(ip_address="120.30.0.1", gateway="192.168.0.1", subnet_mask="255.255.255.0") + node.connect_nic(nic1) + nic2 = NIC(ip_address="120.30.0.2", gateway="192.168.0.1", subnet_mask="255.255.255.0") + node.connect_nic(nic2) + + net.add_node(node) + + with pytest.raises(RuntimeError): + net.connect_nodes(node.nics[nic1.uuid], node.nics[nic2.uuid], bandwidth=30) + + assert node in net + assert nic1.connected_link is None + assert nic2.connected_link is None + assert len(net.links) == 0 + + +def test_disconnecting_nodes(): + ...