diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py index 1feadef1b4..0e27cd5e8c 100644 --- a/python/tests/test_topology.py +++ b/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()