def test_set_jump_existing_trial(self):
lineages = Lineages()
root_1 = TrialStub(id="root-1")
lineage_1 = lineages.add(root_1)
root_2 = TrialStub(id="root-2")
lineage_2 = lineages.add(root_2)
child_trial = TrialStub(id="child")
child_lineage = lineages.fork(root_1, child_trial)
lineages.set_jump(root_2, child_trial)
assert child_lineage.base is lineage_2
assert lineage_2.jumps == [child_lineage]
assert child_lineage.jumps == []
assert lineage_2.base is None
assert lineage_1.jumps == []
assert lineage_1.base is None
def test_set_jump_non_existing_new_trial(self):
lineages = Lineages()
trial = TrialStub(id="exists")
lineages.add(trial)
with pytest.raises(KeyError, match="'newtrialdontexist'"):
lineages.set_jump(trial, TrialStub(id="newtrialdontexist"))
def test_set_jump_non_existing_base_trial(self):
lineages = Lineages()
with pytest.raises(KeyError, match="'dontexist'"):
lineages.set_jump(
TrialStub(id="dontexist"), TrialStub(id="dontexistbutdoesntmatter")
)