Coverage and bugs

talmolab · Apr 14, 2024 · 578f6a4 · 578f6a4
1 parent 6850eb0
commit 578f6a4
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/sleap_io/model/skeleton.py b/sleap_io/model/skeleton.py
@@ -221,8 +221,9 @@ def add_edge(self, src: Edge | Node | str = None, dst: Node | str = None):
             dst: The destination `Node` or name of the destination node.
         """
         if type(src) == Edge:
-            if src not in self.edges:
-                self.edges.append(src)
+            edge = src
+            if edge not in self.edges:
+                self.edges.append(edge)
             if edge.source not in self.nodes:
                 self.add_node(edge.source)
             if edge.destination not in self.nodes:

diff --git a/tests/model/test_skeleton.py b/tests/model/test_skeleton.py
@@ -132,6 +132,10 @@ def test_add_edge():
     skel.add_edge("D", "A")
     assert skel.edge_inds == [(0, 1), (1, 0), (0, 2), (3, 0)]
 
+    skel = Skeleton(["A", "B"])
+    skel.add_edge(Edge(Node("A"), Node("B")))
+    assert skel.edge_inds == [(0, 1)]
+
 
 def test_add_symmetry():
     skel = Skeleton(["A", "B"])
@@ -146,3 +150,10 @@ def test_add_symmetry():
         Symmetry([Node("A"), Node("B")]),
         Symmetry([Node("C"), Node("D")]),
     ]
+
+    skel.add_symmetry("E", "F")
+    assert skel.symmetries == [
+        Symmetry([Node("A"), Node("B")]),
+        Symmetry([Node("C"), Node("D")]),
+        Symmetry([Node("E"), Node("F")]),
+    ]