import unittest

from scripts.snake_arena_scenarios import SCENARIOS, synthetic_states

class TestSnakeArenaScenarios(unittest.TestCase):

  def test_default_corpus_rotates_through_every_scenario(self):
    states = synthetic_states(len(SCENARIOS))

    self.assertEqual(
      {metadata["scenario"] for _, metadata in states},
      set(SCENARIOS),
    )

  def test_scenario_filter_is_deterministic(self):
    first = synthetic_states(3, ["hazard"])
    second = synthetic_states(3, ["hazard"])

    self.assertEqual(first, second)
    self.assertTrue(all(metadata["scenario"] == "hazard" for _, metadata in first))

  def test_generated_you_is_present_on_board(self):
    for board, metadata in synthetic_states(20):
      with self.subTest(scenario=metadata["scenario"]):
        ids = {snake["id"] for snake in board["snakes"]}
        self.assertIn(metadata["you"]["id"], ids)
        self.assertGreater(board["width"], 0)
        self.assertGreater(board["height"], 0)

if __name__ == "__main__":
  unittest.main()