def test_remove(self) -> None: """Test removal of a state from beam.""" beam = Beam(width=2) state1 = State(score=0.0) beam.add_state(state1) beam.remove(state1) state2 = State(score=1.0) beam.add_state(state2) beam.remove(state2) state3 = State(score=2.0) beam.add_state(state3) beam.remove(state3) assert beam.size == 0 beam.add_state(state1) beam.add_state(state2) beam.add_state(state3) assert state1 not in beam.iter_states() assert state2 in beam.iter_states() assert state3 in beam.iter_states() # TODO: uncomment once fixed https://github.com/thoth-station/adviser/issues/1541 # with pytest.raises(ValueError): # beam.remove(state1) assert beam.max() is state3 beam.remove(state2) assert beam.max() is state3 assert state2 not in beam.iter_states() assert state3 in beam.iter_states() assert beam.size == 1 beam.remove(state3) assert beam.size == 0 assert state3 not in beam.iter_states()
def test_remove(self) -> None: """Test removal of a state from beam.""" beam = Beam(width=2) state1 = State(score=0.0) beam.add_state(state1) beam.remove(state1) state2 = State(score=1.0) beam.add_state(state2) beam.remove(state2) state3 = State(score=2.0) beam.add_state(state3) beam.remove(state3) assert beam.size == 0 beam.add_state(state1) beam.add_state(state2) beam.add_state(state3) assert state1 not in beam.iter_states() assert state2 in beam.iter_states() assert state3 in beam.iter_states() with pytest.raises(KeyError): beam.remove(state1) beam.remove(state2) assert state2 not in beam.iter_states() assert state3 in beam.iter_states() assert beam.size == 1 beam.remove(state3) assert beam.size == 0 assert state3 not in beam.iter_states()