one more test

python/tests/test_extend_edges.py

@@ -568,6 +568,41 @@ def test_iterative_example(self):
         assert sts.num_edges == 16
         tables.assert_equals(sts.extend_edges().tables, ignore_provenance=True)
 
+    def test_very_simple(self):
+        samples = [0]
+        node_times = [0, 1, 2, 3]
+        # (p, c, l, r)
+        edges = [
+            (1, 0, 0, 1),
+            (2, 0, 1, 2),
+            (2, 1, 0, 1),
+            (3, 0, 2, 3),
+            (3, 2, 0, 2),
+        ]
+        correct_edges = [
+            (1, 0, 0, 3),
+            (2, 1, 0, 3),
+            (3, 2, 0, 3),
+        ]
+        tables = tskit.TableCollection(sequence_length=3)
+        for n, t in enumerate(node_times):
+            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
+            tables.nodes.add_row(time=t, flags=flags)
+        for p, c, l, r in edges:
+            tables.edges.add_row(parent=p, child=c, left=l, right=r)
+        ts = tables.tree_sequence()
+        ets = ts.extend_edges()
+        for _, t, et in ts.coiterate(ets):
+            print("----")
+            print(t.draw(format="ascii"))
+            print(et.draw(format="ascii"))
+        etables = ets.tables
+        correct_tables = etables.copy()
+        etables.edges.clear()
+        for p, c, l, r in correct_edges:
+            etables.edges.add_row(parent=p, child=c, left=l, right=r)
+        etables.assert_equals(correct_tables, ignore_provenance=True)
+
     def test_wright_fisher(self):
         tables = wf.wf_sim(N=5, ngens=20, num_loci=100, deep_history=False, seed=3)
         tables.sort()