better test

python/tests/test_topology.py

@@ -8191,14 +8191,37 @@ def verify_extend_edges(self, ts, ets):
                     overlaps = True
             assert overlaps
 
-        # TODO: the above is sufficient but not necessary
-        # is there an easy way to check that we're doing
-        # *all* the extending that's possible
-        # and/or that everything we've done is correct?
+        next_chains = []
+        last_chains = []
+        for _, t, et in ts.coiterate(ets):
+            for a in t.nodes():
+                b = t.parent(a)
+                if b != tskit.NULL:
+                    c = t.parent(b)
+                    if c != tskit.NULL:
+                        next_chains.append((a, b, c))
+            for a, b, c in last_chains:
+                if a in t.nodes() and t.parent(a) == c and b not in t.nodes():
+                    # the relationship a <- b <- c should still be in the tree,
+                    # although maybe they aren't direct parent-offspring
+                    assert a in et.nodes()
+                    assert b in et.nodes()
+                    assert c in et.nodes()
+                    p = a
+                    while p != tskit.NULL:
+                        if p == b:
+                            break
+                        p = et.parent(p)
+                    assert p == b
+                    while p != tskit.NULL:
+                        if p == c:
+                            break
+                        p = et.parent(p)
+                    assert p == c
+            last_chains = next_chains
 
-    # TODO: test this on some more tree sequences
     def test_extend_edges(self):
-        tables = wf.wf_sim(10, 1, deep_history=False, seed=1)
+        tables = wf.wf_sim(10, 10, deep_history=False, seed=1)
         tables.sort()
         ts = tables.tree_sequence()
         ets = ts.extend_edges()