def test_top(self) -> None:
"""Test top element in beam."""
beam = Beam(width=2)
assert beam.width == 2
state1 = State(score=1.0)
beam.add_state(state1)
assert beam.top() is state1
state2 = State(score=2.0)
beam.add_state(state2)
assert beam.top() is state2
state3 = State(score=0.5)
beam.add_state(state3)
assert beam.top() is state2
def test_run(self, state_count: int) -> None:
"""Test running the sampling method."""
beam = Beam()
for _ in range(state_count):
state = State(score=random.random())
beam.add_state(state)
predictor = HillClimbing()
context = flexmock(accepted_final_states_count=33)
next_state = predictor.run(context, beam)
assert next_state is not None
assert next_state is beam.top()