refactor: make view.find_element_view and find_relationship_view more generally useful

Author: yt-ms

src/structurizr/view/view.py

@@ -130,7 +130,7 @@ def relationship_views(self) -> Iterable[RelationshipView]:
         """Return the relationship views contained by this view."""
         return set(self._relationship_views)
 
-    def _add_element(self, element: Element, add_relationships: bool) -> None:
+    def _add_element(self, element: Element, add_relationships: bool) -> ElementView:
         """
         Add the given element to this view.
 
@@ -145,10 +145,13 @@ def _add_element(self, element: Element, add_relationships: bool) -> None:
                 f"The element {element} does not exist in the model associated with "
                 f"this view."
             )
-        if element not in [view.element for view in self.element_views]:
-            self.element_views.add(ElementView(element=element))
+        view = self.find_element_view(element=element)
+        if view is None:
+            view = ElementView(element=element)
+            self.element_views.add(view)
         if add_relationships:
             self._add_relationships(element)
+        return view
 
     def _remove_element(self, element: Element) -> None:
         """
@@ -192,15 +195,8 @@ def _add_relationship(
         if self.is_element_in_view(relationship.source) and self.is_element_in_view(
             relationship.destination
         ):
-            view = next(
-                (
-                    rv
-                    for rv in self._relationship_views
-                    if rv.relationship is relationship
-                    and rv.description == description
-                    and rv.response == response
-                ),
-                None,
+            view = self.find_relationship_view(
+                relationship=relationship, description=description, response=response
             )
             if not view:
                 view = RelationshipView(
@@ -211,7 +207,6 @@ def _add_relationship(
                 )
                 self._relationship_views.add(view)
             return view
-        return None
 
     def _add_relationships(self, element: Element) -> None:
         """
@@ -241,15 +236,17 @@ def copy_layout_information_from(self, source: "View") -> None:
             self.paper_size = source.paper_size
 
         for source_element_view in source.element_views:
-            destination_element_view = self.find_element_view(source_element_view)
+            destination_element_view = self.find_element_view(
+                element=source_element_view.element
+            )
             if destination_element_view:
                 destination_element_view.copy_layout_information_from(
                     source_element_view
                 )
 
         for source_relationship_view in source.relationship_views:
             destintion_relationship_view = self.find_relationship_view(
-                source_relationship_view
+                relationship=source_relationship_view.relationship
             )
             if destintion_relationship_view:
                 destintion_relationship_view.copy_layout_information_from(
@@ -258,28 +255,47 @@ def copy_layout_information_from(self, source: "View") -> None:
 
     def is_element_in_view(self, element: Element) -> bool:
         """Return True if the given element is in this view."""
-        return any([e.element.id == element.id for e in self.element_views])
+        return self.find_element_view(element=element) is not None
 
     def find_element_view(
-        self, source_element_view: ElementView
+        self,
+        *,
+        element: Optional[Element] = None,
     ) -> Optional[ElementView]:
-        """Find a child element view corresponding to the given source view."""
-        for element_view in self.element_views:
-            if element_view.element.id == source_element_view.element.id:
-                return element_view
-        return None
+        """Find a child element view matching a given element."""
+        return next(
+            (view for view in self.element_views if view.element.id == element.id), None
+        )
 
     def find_relationship_view(
-        self, source_relationship_view: RelationshipView
+        self,
+        *,
+        relationship: Optional[Relationship] = None,
+        description: Optional[str] = None,
+        response: Optional[bool] = None,
     ) -> Optional[RelationshipView]:
-        """Find a child element view corresponding to the given relationship view."""
-        for relationship_view in self._relationship_views:
+        """
+        Find a child relationship view matching the supplied non-None arguments.
+
+        Args:
+            relationship: find a child view with matching relationship ID
+            description:  find a child view with matching view description.  Note that
+                          the view description is not always the same as that of the
+                          relationship
+            response:     find a child view with matching response indicator.
+        """
+        for view in self._relationship_views:
+            rel = view.relationship
             if (
-                relationship_view.relationship.id
-                == source_relationship_view.relationship.id
+                (relationship is None or rel.id == relationship.id)
+                and (
+                    description is None
+                    or view.description == description
+                    or (view.description is None and rel.description == description)
+                )
+                and (response is None or view.response == response)
             ):
-                return relationship_view
-        return None
+                return view
 
     def check_parent_and_children_not_in_view(self, element: Element) -> None:
         """Ensure that an element can't be added if parent or children are in view."""

tests/unit/view/test_view.py

@@ -16,6 +16,7 @@
 """Ensure the expected behaviour of View."""
 
 from structurizr.model import Model
+from structurizr.view.paper_size import PaperSize
 from structurizr.view.view import View, ViewIO
 
 
@@ -25,6 +26,56 @@ class DerivedView(View):
     pass
 
 
+def test_find_element_view():
+    """Test behaviour of find_element_view."""
+    model = Model()
+    sys1 = model.add_software_system(name="System 1")
+    sys2 = model.add_software_system(name="System 2")
+
+    view = DerivedView(software_system=sys1, description="")
+    view._add_element(sys1, False)
+
+    assert view.find_element_view(element=sys1).element is sys1
+    assert view.find_element_view(element=sys2) is None
+
+
+def test_find_relationship_view():
+    """Test behaviour of find_element_view."""
+    model = Model()
+    sys1 = model.add_software_system(name="System 1")
+    sys2 = model.add_software_system(name="System 2")
+    rel1 = sys1.uses(sys2, "Uses")
+    rel2 = sys2.uses(sys1, "Also uses")
+    rel3 = sys2.uses(sys1, "Returns")
+
+    view = DerivedView(software_system=sys1, description="")
+    view._add_element(sys1, False)
+    view._add_element(sys2, False)
+    view._add_relationship(rel1).description = "Override"
+    view._add_relationship(rel3).response = True
+
+    assert view.find_relationship_view(relationship=rel1).relationship is rel1
+    assert view.find_relationship_view(relationship=rel2) is None
+    assert view.find_relationship_view(description="Override").relationship is rel1
+    assert view.find_relationship_view(description="Uses") is None
+    assert view.find_relationship_view(description="Returns").relationship is rel3
+    assert view.find_relationship_view(response=True).relationship is rel3
+    assert view.find_relationship_view(response=False).relationship is rel1
+
+
+def test_is_element_in_view():
+    """Test check for an element being in the view."""
+    model = Model()
+    sys1 = model.add_software_system(name="System 1")
+    sys2 = model.add_software_system(name="System 2")
+
+    view = DerivedView(software_system=sys1, description="")
+    view._add_element(sys1, False)
+
+    assert view.is_element_in_view(sys1)
+    assert not view.is_element_in_view(sys2)
+
+
 def test_add_relationship_doesnt_duplicate():
     """Test that adding a relationships twice doesn't duplicate it."""
     model = Model()
@@ -97,3 +148,28 @@ def test_missing_json_description_allowed():
     """
     io = ViewIO.parse_raw(json)
     assert io is not None
+
+
+def test_copy_layout():
+    """Ensure that layout is copied over, including sub-views."""
+    model = Model()
+    sys1 = model.add_software_system(name="System 1")
+    sys2 = model.add_software_system(name="System 2")
+    rel1 = sys1.uses(sys2)
+
+    view1 = DerivedView(software_system=sys1, description="")
+    view1._add_element(sys1, False).paper_size = PaperSize.A1_Portrait
+    view1._add_element(sys2, False).paper_size = PaperSize.A2_Portrait
+    view1._add_relationship(rel1).paper_size = PaperSize.A3_Portrait
+    view1.paper_size = PaperSize.A4_Portrait
+
+    view2 = DerivedView(software_system=sys1, description="")
+    view2._add_element(sys1, False).paper_size = PaperSize.A1_Portrait
+    view2._add_element(sys2, False).paper_size = PaperSize.A2_Portrait
+    view2._add_relationship(rel1).paper_size = PaperSize.A3_Portrait
+    view2.copy_layout_information_from(view1)
+
+    assert view2.paper_size == PaperSize.A4_Portrait
+    assert view2.find_element_view(element=sys1).paper_size == PaperSize.A1_Portrait
+    rv = view2.find_relationship_view(description="Uses")
+    assert rv.paper_size == PaperSize.A3_Portrait