feat: add InfrrastructureNodes to DeploymentNodes

Author: yt-ms

src/structurizr/model/deployment_node.py

@@ -201,6 +201,14 @@ def add_software_system_instance(
         model += instance
         return instance
 
+    def add_infrastructure_node(self, name: str, **kwargs) -> InfrastructureNode:
+        """Create a new infrastructure node under this node."""
+        infra_node = InfrastructureNode(name=name, parent=self, **kwargs)
+        self._infrastructure_nodes.add(infra_node)
+        model = self.model
+        model += infra_node
+        return infra_node
+
     def __iadd__(self, node: "DeploymentNode") -> "DeploymentNode":
         """Add a newly constructed chile deployment node to this node."""
         if node in self._children:
@@ -255,4 +263,10 @@ def hydrate(
             )
             node._software_system_instances.add(instance)
 
+        for infra_node_io in deployment_node_io.infrastructure_nodes:
+            infra_node = InfrastructureNode.hydrate(
+                infra_node_io, model=model, parent=node
+            )
+            node._infrastructure_nodes.add(infra_node)
+
         return node

src/structurizr/model/infrastructure_node.py

@@ -20,6 +20,7 @@
 
 
 if TYPE_CHECKING:
+    from .deployment_node import DeploymentNode
     from .model import Model
 
 
@@ -55,18 +56,21 @@ def __init__(
         self,
         *,
         technology: str = "",
+        parent: "DeploymentNode" = None,
         **kwargs,
     ):
         """Initialize an infrastructure node model."""
         super().__init__(**kwargs)
         self.technology = technology
         self.tags.add(Tags.INFRASTRUCTURE_NODE)
+        self.parent = parent
 
     @classmethod
     def hydrate(
         cls,
         node_io: InfrastructureNodeIO,
         model: "Model",
+        parent: "DeploymentNode",
     ) -> "InfrastructureNode":
         """Hydrate a new InfrastructureNode instance from its IO.
 
@@ -75,6 +79,7 @@ def hydrate(
         node = cls(
             **cls.hydrate_arguments(node_io),
             technology=node_io.technology,
+            parent=parent,
         )
         model += node
         return node

tests/unit/model/test_deployment_node.py

@@ -207,3 +207,35 @@ def test_deployment_node_adding_software_system_replicating_relationships(
 ):
     """Test replicating relationships when adding a software system instance."""
     raise AssertionError()  # Not implemented yet
+
+
+def test_deployment_node_add_infrastructure_node(model_with_node):
+    """Test adding an infrastructure node to a deployment node."""
+    node = model_with_node.empty_node
+
+    infra_node = node.add_infrastructure_node("infraNode")
+
+    assert infra_node.name == "infraNode"
+    assert infra_node.model is model_with_node
+    assert infra_node.parent is node
+    assert infra_node in node.infrastructure_nodes
+
+
+def test_deployment_node_serialising_infrastructure_nodes(model_with_node):
+    """Test serialisation and deserialisation includes infrastructure nodes."""
+    node = model_with_node.empty_node
+    node.add_infrastructure_node("infraNode")
+
+    io = DeploymentNodeIO.from_orm(node)
+
+    assert len(io.infrastructure_nodes) == 1
+    assert io.infrastructure_nodes[0].id == "id"
+    assert io.infrastructure_nodes[0].name == "infraNode"
+
+    node2 = DeploymentNode.hydrate(io, model_with_node)
+
+    assert len(node2.infrastructure_nodes) == 1
+    infra_node = node2.infrastructure_nodes[0]
+    assert infra_node.name == "infraNode"
+    assert infra_node.model is model_with_node
+    assert infra_node.parent is node2

tests/unit/model/test_infrastructure_node.py

@@ -68,8 +68,10 @@ def test_infrastructure_node_tags():
 def test_infrastructure_node_hydration(empty_model):
     """Check hydrating an infrastructure node from its IO."""
     io = InfrastructureNodeIO(name="node1", technology="tech")
+    parent = object()
 
-    node = InfrastructureNode.hydrate(io, empty_model)
+    node = InfrastructureNode.hydrate(io, empty_model, parent=parent)
 
     assert node.name == "node1"
     assert node.technology == "tech"
+    assert node.parent is parent