feat(DeploymentNode): support adding infra nodes with +=

Author: yt-ms

src/structurizr/model/deployment_node.py

@@ -226,23 +226,30 @@ def add_software_system(
     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
+        self._add_infrastructure_node(infra_node)
         return infra_node
 
     def __iadd__(
-        self, child: Union["DeploymentNode", Container, SoftwareSystem]
+        self,
+        child: Union[Container, "DeploymentNode", InfrastructureNode, SoftwareSystem],
     ) -> "DeploymentNode":
-        """Add a sub-node, container or system to this node."""
+        """Add a sub-node, container, system or infra node to this node."""
         if isinstance(child, SoftwareSystem):
             self.add_software_system(child)
         elif isinstance(child, Container):
             self.add_container(child)
+        elif isinstance(child, InfrastructureNode):
+            self._add_infrastructure_node(child)
         else:
             self._add_child_deployment_node(child)
         return self
 
+    def _add_infrastructure_node(self, infra_node: InfrastructureNode):
+        """Add a new infrastructure node."""
+        self._infrastructure_nodes.add(infra_node)
+        model = self.model
+        model += infra_node
+
     def _add_child_deployment_node(self, node: "DeploymentNode"):
         """Add a newly constructed child deployment node to this node."""
         if node in self._children:

tests/unit/model/test_deployment_node.py

@@ -18,6 +18,7 @@
 
 from structurizr.model import Container, SoftwareSystem
 from structurizr.model.deployment_node import DeploymentNode, DeploymentNodeIO
+from structurizr.model.infrastructure_node import InfrastructureNode
 
 
 class MockModel:
@@ -153,6 +154,7 @@ def test_deployment_node_add_with_iadd(model_with_node: MockModel):
     container = Container(name="container")
     system += container
     child_node = DeploymentNode(name="child")
+    infra_node = InfrastructureNode(name="infra")
 
     node += child_node
     assert child_node in node.children
@@ -163,6 +165,9 @@ def test_deployment_node_add_with_iadd(model_with_node: MockModel):
     child_node += container
     assert child_node.container_instances[0].container is container
 
+    child_node += infra_node
+    assert infra_node in child_node.infrastructure_nodes
+
 
 def test_deployment_node_serialising_container(model_with_node):
     """Test serialisation and deserialisation includes container instances."""