From b4cc1b43790abb09eb0995b177c8451800e7c3dd Mon Sep 17 00:00:00 2001
From: Marek Wolan <marek.wolan@methods.co.uk>
Date: Fri, 4 Oct 2024 11:07:49 +0100
Subject: [PATCH] Add tests for office lan creation

---
 src/primaite/simulator/network/creation.py    |  4 +-
 .../_simulator/_network/test_creation.py      | 45 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit_tests/_primaite/_simulator/_network/test_creation.py

diff --git a/src/primaite/simulator/network/creation.py b/src/primaite/simulator/network/creation.py
index 2eefa75f..80d0c370 100644
--- a/src/primaite/simulator/network/creation.py
+++ b/src/primaite/simulator/network/creation.py
@@ -1,7 +1,7 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 from abc import ABC, abstractmethod
 from ipaddress import IPv4Address
-from typing import Any, ClassVar, Dict, Type
+from typing import Any, ClassVar, Dict, Literal, Type
 
 from pydantic import BaseModel
 
@@ -98,6 +98,7 @@ class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"):
     class ConfigSchema(NetworkNodeAdder.ConfigSchema):
         """Configuration schema for OfficeLANAdder."""
 
+        type: Literal["office_lan"] = "office_lan"
         lan_name: str
         """Name of lan used for generating hostnames for new nodes."""
         subnet_base: int
@@ -197,6 +198,7 @@ class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"):
             # Create and add a PC to the network
             pc = Computer(
                 hostname=f"pc_{i}_{config.lan_name}",
+                # TODO: what happens when ip_block_start + num_pcs exceeds 254?
                 ip_address=f"192.168.{config.subnet_base}.{i+config.pcs_ip_block_start-1}",
                 subnet_mask="255.255.255.0",
                 default_gateway=default_gateway,
diff --git a/tests/unit_tests/_primaite/_simulator/_network/test_creation.py b/tests/unit_tests/_primaite/_simulator/_network/test_creation.py
new file mode 100644
index 00000000..f5b19179
--- /dev/null
+++ b/tests/unit_tests/_primaite/_simulator/_network/test_creation.py
@@ -0,0 +1,45 @@
+# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
+from ipaddress import IPv4Address
+
+import pytest
+
+from primaite.simulator.network.container import Network
+from primaite.simulator.network.creation import OfficeLANAdder
+
+
+@pytest.mark.parametrize(
+    ("lan_name", "subnet_base", "pcs_ip_block_start", "num_pcs", "include_router", "bandwidth"),
+    (
+        ("CORP-NETWORK", 3, 10, 6, True, 45),
+        ("OTHER-NETWORK", 10, 25, 26, True, 100),
+        ("OTHER-NETWORK", 10, 25, 55, False, 100),
+    ),
+)
+def test_office_lan_adder(lan_name, subnet_base, pcs_ip_block_start, num_pcs, include_router, bandwidth):
+    net = Network()
+
+    office_lan_config = OfficeLANAdder.ConfigSchema(
+        lan_name=lan_name,
+        subnet_base=subnet_base,
+        pcs_ip_block_start=pcs_ip_block_start,
+        num_pcs=num_pcs,
+        include_router=include_router,
+        bandwidth=bandwidth,
+    )
+    OfficeLANAdder.add_nodes_to_net(config=office_lan_config, network=net)
+
+    num_switches = 1 if num_pcs <= 23 else num_pcs // 23 + 2
+    num_routers = 1 if include_router else 0
+    total_nodes = num_pcs + num_switches + num_routers
+
+    assert all((n.hostname.endswith(lan_name) for n in net.nodes.values()))
+    assert len(net.computer_nodes) == num_pcs
+    assert len(net.switch_nodes) == num_switches
+    assert len(net.router_nodes) == num_routers
+    assert len(net.nodes) == total_nodes
+    assert all(
+        [str(n.network_interface[1].ip_address).startswith(f"192.168.{subnet_base}") for n in net.computer_nodes]
+    )
+    # check that computers occupy address range 192.168.3.10 - 192.168.3.16
+    computer_ip_last_octets = {str(n.network_interface[1].ip_address).split(".")[-1] for n in net.computer_nodes}
+    assert computer_ip_last_octets == {str(i) for i in range(pcs_ip_block_start, pcs_ip_block_start + num_pcs)}