feat(DeploymentNode): add containers and systems with +=

Author: yt-ms

src/structurizr/model/deployment_node.py

@@ -15,7 +15,7 @@
 
 """Provide a deployment node model."""
 
-from typing import TYPE_CHECKING, Iterable, List
+from typing import TYPE_CHECKING, Iterable, List, Union
 
 from pydantic import Field
 
@@ -148,7 +148,7 @@ def add_deployment_node(
         node = DeploymentNode(
             name=name, description=description, technology=technology, **kwargs
         )
-        self += node
+        self._add_child_deployment_node(node)
         return node
 
     def add_container(
@@ -231,8 +231,20 @@ def add_infrastructure_node(self, name: str, **kwargs) -> InfrastructureNode:
         model += infra_node
         return infra_node
 
-    def __iadd__(self, node: "DeploymentNode") -> "DeploymentNode":
-        """Add a newly constructed chile deployment node to this node."""
+    def __iadd__(
+        self, child: Union["DeploymentNode", Container, SoftwareSystem]
+    ) -> "DeploymentNode":
+        """Add a sub-node, container or system to this node."""
+        if isinstance(child, SoftwareSystem):
+            self.add_software_system(child)
+        elif isinstance(child, Container):
+            self.add_container(child)
+        else:
+            self._add_child_deployment_node(child)
+        return self
+
+    def _add_child_deployment_node(self, node: "DeploymentNode"):
+        """Add a newly constructed child deployment node to this node."""
         if node in self._children:
             return self
 
@@ -252,7 +264,6 @@ def __iadd__(self, node: "DeploymentNode") -> "DeploymentNode":
         self._children.add(node)
         model = self.model
         model += node
-        return self
 
     @classmethod
     def hydrate(

tests/unit/model/test_deployment_node.py

@@ -16,6 +16,7 @@
 
 import pytest
 
+from structurizr.model import Container, SoftwareSystem
 from structurizr.model.deployment_node import DeploymentNode, DeploymentNodeIO
 
 
@@ -40,6 +41,10 @@ def get_element(self, id: str):
         assert id == self.mock_element.id
         return self.mock_element
 
+    def get_elements(self):
+        """Simulate get_elements."""
+        return []
+
 
 class MockElement:
     """Implement a mock element for testing."""
@@ -140,6 +145,25 @@ def test_deployment_node_add_container(model_with_node):
     assert instance.instance_id == 1
 
 
+def test_deployment_node_add_with_iadd(model_with_node: MockModel):
+    """Test adding things to a node using += rather than add_container."""
+    node = model_with_node.empty_node
+    system = SoftwareSystem(name="system")
+    model_with_node += system
+    container = Container(name="container")
+    system += container
+    child_node = DeploymentNode(name="child")
+
+    node += child_node
+    assert child_node in node.children
+
+    node += system
+    assert node.software_system_instances[0].software_system is system
+
+    child_node += container
+    assert child_node.container_instances[0].container is container
+
+
 def test_deployment_node_serialising_container(model_with_node):
     """Test serialisation and deserialisation includes container instances."""
     node = model_with_node.empty_node