def test_cpuagent_get_state_correctness():
agent = CPUAgent()
agent.states = [
State(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])),
State(np.array([0, 0, 0, 0, 0, 0, 0, 0, 1])),
]
result = agent.get_state(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])).grid
assert_array_equal(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]), result)
def test_cpuagent_get_state_raises_value_error_if_state_not_found():
agent = CPUAgent()
with pytest.raises(ValueError, match="could not be found in saved states"):
agent.get_state(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0]))