style: PR feedback

Author: yt-ms

src/structurizr/model/element.py

@@ -102,10 +102,16 @@ def get_afferent_relationships(self) -> Iterator[Relationship]:
     def add_relationship(
         self, relationship: Optional[Relationship] = None, **kwargs
     ) -> Relationship:
+        """Add a new relationship from this element to another.
+
+        This can be used either to add a `Relationship` instance that has already been constructed,
+        or by passing the arguments (e.g. description, destination) with which to construct a new
+        one.  The relationship will automatically be registered with the element's model.
+        """
         if relationship is None:
             relationship = Relationship(**kwargs)
         elif relationship in self.relationships:
-            return relationship  # Nothing more to do
+            return relationship
         if relationship.source is None:
             relationship.source = self
         elif relationship.source is not self:

src/structurizr/model/model.py

@@ -411,7 +411,7 @@ def _add_element(self, element: Element) -> None:
     def _add_relationship(self, relationship: Relationship) -> bool:
         """"""
         if relationship in self.get_relationships():
-            return True  # Nothing to do
+            return True
         if not relationship.id:
             relationship.id = self._id_generator.generate_id()
         elif relationship.id in self._elements_by_id:

tests/integration/test_model_element_relationships.py

@@ -36,12 +36,14 @@ def test_adding_relationship_to_element_adds_to_model():
     sys1 = model.add_software_system(name="sys1")
     sys2 = model.add_software_system(name="sys2")
 
-    r = sys1.add_relationship(source=sys1, destination=sys2, description="uses")
-    assert list(sys1.relationships) == [r]
-    assert list(sys1.get_relationships()) == [r]
-    assert list(model.get_relationships()) == [r]
-    assert list(sys2.relationships) == []  # relationships only contains outbound
-    assert list(sys2.get_relationships()) == [r]
+    relationship = sys1.add_relationship(
+        source=sys1, destination=sys2, description="uses"
+    )
+    assert sys1.relationships == {relationship}
+    assert set(sys1.get_relationships()) == {relationship}
+    assert set(model.get_relationships()) == {relationship}
+    assert sys2.relationships == set()  # relationships only contains outbound
+    assert set(sys2.get_relationships()) == {relationship}
 
 
 def test_adding_relationship_to_model_adds_to_element():
@@ -53,12 +55,14 @@ def test_adding_relationship_to_model_adds_to_element():
     sys1 = model.add_software_system(name="sys1")
     sys2 = model.add_software_system(name="sys2")
 
-    r = model.add_relationship(source=sys1, destination=sys2, description="uses")
-    assert list(sys1.relationships) == [r]
-    assert list(sys1.get_relationships()) == [r]
-    assert list(model.get_relationships()) == [r]
-    assert list(sys2.relationships) == []  # relationships only contains outbound
-    assert list(sys2.get_relationships()) == [r]
+    relationship = model.add_relationship(
+        source=sys1, destination=sys2, description="uses"
+    )
+    assert sys1.relationships == {relationship}
+    assert set(sys1.get_relationships()) == {relationship}
+    assert set(model.get_relationships()) == {relationship}
+    assert sys2.relationships == set()  # relationships only contains outbound
+    assert set(sys2.get_relationships()) == {relationship}
 
 
 def test_adding_relationship_via_uses_adds_to_elements():
@@ -70,12 +74,12 @@ def test_adding_relationship_via_uses_adds_to_elements():
     sys1 = model.add_software_system(name="sys1")
     sys2 = model.add_software_system(name="sys2")
 
-    r = sys1.uses(sys2, "uses")
-    assert list(sys1.relationships) == [r]
-    assert list(sys1.get_relationships()) == [r]
-    assert list(model.get_relationships()) == [r]
-    assert list(sys2.relationships) == []  # relationships only contains outbound
-    assert list(sys2.get_relationships()) == [r]
+    relationship = sys1.uses(sys2, "uses")
+    assert sys1.relationships == {relationship}
+    assert set(sys1.get_relationships()) == {relationship}
+    assert set(model.get_relationships()) == {relationship}
+    assert sys2.relationships == set()  # relationships only contains outbound
+    assert set(sys2.get_relationships()) == {relationship}
 
 
 @pytest.mark.parametrize(

tests/unit/model/test_element.py

@@ -19,7 +19,6 @@
 import pytest
 
 from structurizr.model.element import Element
-from structurizr.model.model import Model
 
 
 class ConcreteElement(Element):
@@ -31,7 +30,8 @@ class ConcreteElement(Element):
 class MockModel:
     """Implement a mock model for reference testing."""
 
-    pass
+    def add_relationship(self, relationship):
+        return relationship
 
 
 @pytest.mark.parametrize(
@@ -74,10 +74,10 @@ def test_element_add_relationship_can_omit_source():
     """
     element1 = ConcreteElement(name="elt1")
     element2 = ConcreteElement(name="elt1")
-    model = Model()
-    model._add_element(element1)
-    r = element1.add_relationship(destination=element2)
-    assert r.source is element1
+    model = MockModel()
+    element1.set_model(model)
+    relationship = element1.add_relationship(destination=element2)
+    assert relationship.source is element1
 
 
 def test_element_add_relationship_twice_is_ok():
@@ -86,8 +86,8 @@ def test_element_add_relationship_twice_is_ok():
     """
     element1 = ConcreteElement(name="elt1")
     element2 = ConcreteElement(name="elt1")
-    model = Model()
-    model._add_element(element1)
-    r = element1.add_relationship(destination=element2)
-    element1.add_relationship(r)
-    assert list(element1.relationships) == [r]
+    model = MockModel()
+    element1.set_model(model)
+    relationship = element1.add_relationship(destination=element2)
+    element1.add_relationship(relationship)
+    assert element1.relationships == {relationship}

tests/unit/model/test_model.py

@@ -29,29 +29,31 @@ def test_model_get_relationship_by_id(empty_model: Model):
     """Test retrieving relationships by their IDs."""
     sys1 = empty_model.add_software_system(name="sys1")
     sys2 = empty_model.add_software_system(name="sys2")
-    r = empty_model.add_relationship(source=sys1, destination=sys2, id="r1")
-    assert empty_model.get_relationship("r1") is r
+    relationship = empty_model.add_relationship(source=sys1, destination=sys2, id="r1")
+    assert empty_model.get_relationship("r1") is relationship
 
 
 def test_model_add_relationship_twice_ignored(empty_model: Model):
     """Make sure that adding an existing relationship to the Model makes no difference."""
     sys1 = empty_model.add_software_system(name="sys1")
     sys2 = empty_model.add_software_system(name="sys2")
-    r = empty_model.add_relationship(source=sys1, destination=sys2)
-    assert list(empty_model.get_relationships()) == [r]
-    empty_model.add_relationship(r)
-    assert list(empty_model.get_relationships()) == [r]
+    relationship = empty_model.add_relationship(source=sys1, destination=sys2)
+    assert set(empty_model.get_relationships()) == {relationship}
+    empty_model.add_relationship(relationship)
+    assert set(empty_model.get_relationships()) == {relationship}
 
 
 def test_model_cannot_add_relationship_with_same_id_as_existing(empty_model: Model):
     """Ensure you can't add two relationships with the same ID."""
     sys1 = empty_model.add_software_system(name="sys1")
     sys2 = empty_model.add_software_system(name="sys2")
-    r1 = empty_model.add_relationship(source=sys1, destination=sys2, id="r1")
+    relationship1 = empty_model.add_relationship(source=sys1, destination=sys2, id="r1")
     with pytest.raises(
         ValueError, match="Relationship.* has the same ID as Relationship.*"
     ):
-        r2 = empty_model.add_relationship(source=sys1, destination=sys2, id="r1")
+        relationship2 = empty_model.add_relationship(
+            source=sys1, destination=sys2, id="r1"
+        )
 
 
 def test_model_cannot_add_relationship_with_same_id_as_element(empty_model: Model):