def test_get_lineage_existing_node_trial(self):
lineages = Lineages()
for root_index in range(2):
trial = TrialStub(id=f"lineage-{root_index}-0")
lineage = lineages.add(trial)
for depth in range(1, 10):
new_trial = TrialStub(id=f"lineage-{root_index}-{depth}")
lineage = lineages.fork(trial, new_trial)
trial = new_trial
lineage = lineages.get_lineage(TrialStub(id="lineage-0-2"))
assert lineage.root is lineages._lineage_roots[0]
assert lineage.node_depth == 2
lineage = lineages.get_lineage(TrialStub(id="lineage-1-5"))
assert lineage.root is lineages._lineage_roots[1]
assert lineage.node_depth == 5
def test_get_lineage_non_existing_trial(self):
lineages = Lineages()
with pytest.raises(KeyError):
lineages.get_lineage(TrialStub(id="id"))
def test_get_lineage_existing_root_trial(self):
lineages = Lineages()
trial = TrialStub(id="stub")
lineage = lineages.add(trial)
assert lineages.get_lineage(trial) is lineage