def test_add_state_order_single(self) -> None:
"""Test adding states to beam and order during addition when score is same - iteration is relevant."""
beam = Beam(width=1)
state01 = State(
score=0.0,
latest_version_offset=0,
iteration=1,
iteration_states_added=0,
)
state01.add_justification([{"state": "01"}])
beam.add_state(state01)
state02 = State(
score=0.0,
latest_version_offset=0,
iteration=0,
iteration_states_added=0,
)
state02.add_justification([{"state": "02"}])
beam.add_state(state02)
assert list(beam.iter_states()) == [state01]
assert list(beam.iter_states_sorted(reverse=True)) == [state01]
assert list(beam.iter_states_sorted(reverse=False)) == [state01]
def test_add_state_order_multi(self) -> None:
"""Test adding states to beam and order during addition when score is same."""
beam = Beam(width=2)
state01 = State(score=0.0)
state01.add_justification(self.JUSTIFICATION_SAMPLE_1)
beam.add_state(state01)
state02 = State(score=0.0)
state02.add_justification(self.JUSTIFICATION_SAMPLE_2)
beam.add_state(state02)
state03 = State(score=0.0)
state03.add_justification(self.JUSTIFICATION_SAMPLE_3)
beam.add_state(state03)
assert list(beam.iter_states_sorted()) == [state01, state02]
assert list(beam.iter_states_sorted(reverse=True)) == [state01, state02]
assert list(beam.iter_states_sorted(reverse=False)) == [state01, state02]
def test_add_state_order_single(self) -> None:
"""Test adding states to beam and order during addition when score is same - iteration is relevant."""
beam = Beam(width=1)
state01 = State(
score=0.0,
iteration=1,
)
state01.add_justification(self.JUSTIFICATION_SAMPLE_1)
beam.add_state(state01)
state02 = State(
score=0.0,
iteration=0,
)
state02.add_justification(self.JUSTIFICATION_SAMPLE_2)
beam.add_state(state02)
assert list(beam.iter_states()) == [state01]
assert list(beam.iter_states_sorted(reverse=True)) == [state01]
assert list(beam.iter_states_sorted(reverse=False)) == [state01]
def test_wipe(self) -> None:
"""Test wiping out beam states."""
beam = Beam()
state1 = State(score=1.0)
beam.add_state(state1)
state2 = State(score=0.0)
beam.add_state(state2)
assert list(beam.iter_states_sorted()) == [state1, state2]
assert list(beam.iter_states()) == [state2, state1]
assert beam.wipe() is None
assert list(beam.iter_states_sorted()) == []
assert list(beam.iter_states()) == []
beam.add_state(state1)
beam.add_state(state2)
assert list(beam.iter_states_sorted()) == [state1, state2]
assert list(beam.iter_states()) == [state2, state1]
def test_add_state_order_multi(self) -> None:
"""Test adding states to beam and order during addition when score is same."""
beam = Beam(width=2)
state01 = State(score=0.0, iteration=0, iteration_states_added=0)
state01.add_justification([{"state": "01"}])
beam.add_state(state01)
state02 = State(score=0.0, iteration=0, iteration_states_added=1)
state02.add_justification([{"state": "02"}])
beam.add_state(state02)
state03 = State(score=0.0, iteration=1, iteration_states_added=0)
state03.add_justification([{"state": "03"}])
beam.add_state(state03)
state04 = State(score=0.0, iteration=1, iteration_states_added=1)
state04.add_justification([{"state": "04"}])
beam.add_state(state04)
assert list(beam.iter_states_sorted()) == [state03, state04]
assert list(beam.iter_states_sorted(reverse=True)) == [state03, state04]
assert list(beam.iter_states_sorted(reverse=False)) == [state04, state03]
def test_iter_states_sorted(self) -> None:
"""Test asking for states returns a sorted list of states."""
beam = Beam(width=4)
assert beam.width == 4
state1 = State(score=1.0)
beam.add_state(state1)
state3 = State(score=3.0)
beam.add_state(state3)
state0 = State(score=0.0)
beam.add_state(state0)
state2 = State(score=2.0)
beam.add_state(state2)
assert list(beam.iter_states_sorted()) == [state3, state2, state1, state0]