snake-python/tests/test_GameplaySchema.py

import sqlite3
import tempfile
import unittest
from contextlib import closing
from pathlib import Path

from server.database.backend import postgresql_schema
from server.database.backend.PostgresqlGameplayBackend import PostgresqlGameplayBackend
from server.database.backend.SqliteGameplayBackend import SqliteGameplayBackend

REDUNDANT = {"idx_turns_game_turn", "idx_snake_turns_game_turn"}

class _RecordingConnection:
  def __init__(self) -> None:
    self.statements:list[str] = []

  async def execute(self, statement:str, *args) -> str:
    self.statements.append(statement)
    return "OK"

class _AcquireContext:
  def __init__(self, connection) -> None:
    self._connection = connection

  async def __aenter__(self):
    return self._connection

  async def __aexit__(self, *exc) -> None:
    return None

class _Pool:
  def __init__(self, connection) -> None:
    self._connection = connection

  def acquire(self):
    return _AcquireContext(self._connection)

class TestPostgresqlSchema(unittest.IsolatedAsyncioTestCase):
  def test_ddl_defines_compact_tables_without_legacy_indexes(self):
    ddl = postgresql_schema.TABLES_DDL
    for table in ("game_turns", "snake_states", "game_replays"):
      self.assertIn(f"CREATE TABLE IF NOT EXISTS {table}", ddl)
    for name in REDUNDANT:
      self.assertNotIn(f"CREATE INDEX IF NOT EXISTS {name}", ddl)
    self.assertIn("SET STORAGE EXTERNAL", ddl)

  async def test_schema_init_is_locked_and_drops_redundant_indexes_separately(self):
    connection = _RecordingConnection()
    backend = PostgresqlGameplayBackend("postgresql://example")
    backend._pool = _Pool(connection)

    await backend._initialize_schema()

    self.assertIn("pg_advisory_lock", connection.statements[0])
    drops = [s for s in connection.statements if "DROP INDEX" in s]
    self.assertEqual({s.split()[-1] for s in drops}, REDUNDANT)
    for statement in drops:
      # CONCURRENTLY only works as a standalone statement outside a transaction.
      self.assertIn("CONCURRENTLY IF EXISTS", statement)
      self.assertNotIn(";", statement)

class TestSqliteSchema(unittest.TestCase):
  def _objects(self, db_path:Path, kind:str) -> set[str]:
    with closing(sqlite3.connect(db_path)) as connection:
      return {row[0] for row in connection.execute("SELECT name FROM sqlite_master WHERE type = ?", (kind,))}

  def test_new_database_has_compact_tables(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      db_path = Path(temp_dir) / "gameplay.sqlite3"
      SqliteGameplayBackend(str(db_path))
      tables = self._objects(db_path, "table")
      self.assertTrue({"games", "game_snakes", "game_turns", "snake_states", "game_replays"} <= tables)
      self.assertFalse({"turns", "snake_turns"} & tables)
      self.assertFalse(self._objects(db_path, "index") & REDUNDANT)

  def test_legacy_database_gets_ids_and_keeps_legacy_tables(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      db_path = Path(temp_dir) / "gameplay.sqlite3"
      with closing(sqlite3.connect(db_path)) as connection, connection:
        connection.executescript("""
          CREATE TABLE games (game_id TEXT PRIMARY KEY, started_at TEXT NOT NULL,
            winner_you INTEGER NOT NULL DEFAULT 0, final_turn INTEGER NOT NULL DEFAULT 0,
            status TEXT NOT NULL DEFAULT 'running', width INTEGER, height INTEGER,
            source TEXT, map_name TEXT, ruleset_name TEXT, ruleset_version TEXT,
            your_snake_id TEXT, your_snake_name TEXT);
          CREATE TABLE game_snakes (game_id TEXT NOT NULL, snake_id TEXT NOT NULL,
            snake_name TEXT, is_you INTEGER NOT NULL DEFAULT 0, PRIMARY KEY (game_id, snake_id));
          CREATE TABLE turns (game_id TEXT, turn INTEGER);
          CREATE TABLE snake_turns (game_id TEXT, turn INTEGER, snake_id TEXT);
          CREATE INDEX idx_turns_game_turn ON turns(game_id, turn);
          INSERT INTO games (game_id, started_at) VALUES ('a', '2026-01-01'), ('b', '2026-01-02');
          INSERT INTO game_snakes (game_id, snake_id) VALUES ('a', 's1'), ('a', 's2');
        """)

      SqliteGameplayBackend(str(db_path))

      with closing(sqlite3.connect(db_path)) as connection:
        gids = [row[0] for row in connection.execute("SELECT gid FROM games ORDER BY game_id")]
        gsids = [row[0] for row in connection.execute("SELECT gsid FROM game_snakes ORDER BY snake_id")]
      self.assertEqual(len(set(gids)), 2)
      self.assertNotIn(None, gids)
      self.assertEqual(len(set(gsids)), 2)
      self.assertNotIn(None, gsids)
      self.assertTrue({"turns", "snake_turns"} <= self._objects(db_path, "table"))
      self.assertFalse(self._objects(db_path, "index") & REDUNDANT)

  def test_running_game_lookups_use_primary_keys(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      db_path = Path(temp_dir) / "gameplay.sqlite3"
      SqliteGameplayBackend(str(db_path))
      with closing(sqlite3.connect(db_path)) as connection:
        for query in (
          "SELECT * FROM game_turns WHERE gid = 1 ORDER BY turn",
          "SELECT * FROM snake_states WHERE gsid = 1 AND turn = 3",
        ):
          plan = " ".join(str(row[-1]) for row in connection.execute(f"EXPLAIN QUERY PLAN {query}"))
          self.assertIn("PRIMARY KEY", plan, plan)

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