feat: minor improvements to Element.add_relationship()

Author: yt-ms

src/structurizr/model/element.py

@@ -72,6 +72,7 @@ def __init__(
         self.name = name
         self.description = description
         self.url = url
+        # Note: relationships should always match get_efferent_relationships() - i.e. outbound relationships only
         self.relationships: Iterable[Relationship] = set(relationships)
 
         self.tags.add(Tags.ELEMENT)
@@ -98,10 +99,15 @@ def get_afferent_relationships(self) -> Iterator[Relationship]:
             r for r in self.get_model().get_relationships() if self is r.destination
         )
 
-    def add_relationship(self, relationship: Optional[Relationship] = None, **kwargs):
+    def add_relationship(self, relationship: Optional[Relationship] = None, **kwargs) -> Relationship:
         if relationship is None:
             relationship = Relationship(**kwargs)
+        if relationship.source is None:
+            relationship.source = self
+        elif relationship.source is not self:
+            raise ValueError(f"Cannot add relationship {relationship} to element {self} that is not its source.")
         self.relationships.add(relationship)
+        return relationship
 
     @classmethod
     def hydrate_arguments(cls, element_io: ElementIO) -> dict:

tests/unit/model/test_element.py

@@ -53,3 +53,24 @@ def test_model_reference():
     element = ConcreteElement(name="Element")
     element.set_model(model)
     assert element.get_model() is model
+
+
+def test_element_can_only_add_relationship_to_source():
+    """Make sure that nothing adds a relationship to the wrong element."""
+    element1 = ConcreteElement(name="elt1")
+    element2 = ConcreteElement(name="elt1")
+    with pytest.raises(
+        ValueError, match="Cannot add relationship .* to element .* that is not its source"
+    ):
+        element1.add_relationship(source=element2)
+
+
+def test_element_add_relationship_can_omit_source():
+    """
+    When using add_relationship with params rather than an existing instance,
+    it should default the source.
+    """
+    element1 = ConcreteElement(name="elt1")
+    element2 = ConcreteElement(name="elt1")
+    r = element1.add_relationship(destination=element2)
+    assert r.source is element1