Extends D2HashableConnection's __eq__ method to also validate label

netsatan · Apr 29, 2024 · cf99f19 · cf99f19
1 parent 1a7800c
commit cf99f19
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/specifipy/diagram_engines/hashable_connection.py b/specifipy/diagram_engines/hashable_connection.py
@@ -6,10 +6,11 @@ def __hash__(self):
         return hash((self.shape_1, self.shape_2, self.direction))
 
     def __eq__(self, other) -> bool:
-        if (self.shape_1, self.shape_2, self.direction) == (
+        if (self.shape_1, self.shape_2, self.direction, self.label) == (
             other.shape_1,
             other.shape_2,
             other.direction,
+            other.label,
         ):
             return True
         return False
diff --git a/tests/test_diagram_engines.py b/tests/test_diagram_engines.py
@@ -10,3 +10,18 @@ def test_connections_uniqueness(self):
         connection_1 = D2HashableConnection("a", "b", "some_label", Direction.TO)
         connection_2 = D2HashableConnection("a", "b", "some_label", Direction.TO)
         assert connection_1 == connection_2
+
+    def test_connections_hash(self):
+        connection_1 = D2HashableConnection("a", "b", "some_label", Direction.TO)
+        assert isinstance(connection_1.__hash__(), int)
+
+    def test_connections_not_equal(self):
+        connection_1 = D2HashableConnection("a", "b", "some_label", Direction.TO)
+        connection_2 = D2HashableConnection("a", "b", "some_label2", Direction.TO)
+        connection_3 = D2HashableConnection("a1", "b", "some_label", Direction.TO)
+        connection_4 = D2HashableConnection("a", "b1", "some_label", Direction.TO)
+        connection_5 = D2HashableConnection("a", "b", "some_label", Direction.FROM)
+        assert connection_1 != connection_2
+        assert connection_1 != connection_3
+        assert connection_1 != connection_4
+        assert connection_1 != connection_5