snake-python/tests/test_MigrateGameplayDatabase.py

import sqlite3
import tempfile
import unittest
from pathlib import Path

from scripts.migrate_gameplay_database import copy_game_snakes
from server.database.backend.SqliteGameplayBackend import SqliteGameplayBackend

class TestMigrateGameplayDatabase(unittest.TestCase):
  def test_copy_game_snakes_preserves_customizations(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      root = Path(temp_dir)
      source_path = root / "source.sqlite3"
      destination_path = root / "destination.sqlite3"
      SqliteGameplayBackend(str(source_path))
      SqliteGameplayBackend(str(destination_path))

      with sqlite3.connect(source_path) as source:
        source.execute("PRAGMA foreign_keys = OFF")
        source.execute("""
          INSERT INTO game_snakes (
            game_id, snake_id, snake_name, is_you, customizations_json
          ) VALUES (?, ?, ?, ?, ?)
        """, (
          "game-1", "snake-1", "PrismBattleSnake", 1,
          '{"color":"#663399","head":"ferret","tail":"swirl"}',
        ))

      source = sqlite3.connect(source_path)
      destination = sqlite3.connect(destination_path)
      try:
        copied = copy_game_snakes(
          source, destination, batch_size=10, retained_ids={"game-1"},
        )
        destination.commit()
        row = destination.execute("""
          SELECT snake_name, is_you, customizations_json
          FROM game_snakes WHERE game_id = ? AND snake_id = ?
        """, ("game-1", "snake-1")).fetchone()
      finally:
        source.close()
        destination.close()

      self.assertEqual(copied, 1)
      self.assertEqual(row, (
        "PrismBattleSnake", 1,
        '{"color":"#663399","head":"ferret","tail":"swirl"}',
      ))

  def test_copy_game_snakes_synthesizes_rows_when_table_is_empty(self):
    source = sqlite3.connect(":memory:")
    destination = sqlite3.connect(":memory:")
    try:
      source.execute("""
        CREATE TABLE game_snakes (
          game_id TEXT, snake_id TEXT, snake_name TEXT, is_you INTEGER,
          customizations_json TEXT NOT NULL DEFAULT '{}'
        )
      """)
      source.execute("""
        CREATE TABLE snake_turns (
          game_id TEXT, snake_id TEXT, snake_name TEXT, is_you INTEGER
        )
      """)
      source.executemany(
        "INSERT INTO snake_turns VALUES (?, ?, ?, ?)",
        [
          ("game-1", "snake-1", "PrismBattleSnake", 1),
          ("game-1", "snake-1", "PrismBattleSnake", 1),
          ("game-1", "snake-2", "Enemy", 0),
        ],
      )
      destination.execute("""
        CREATE TABLE game_snakes (
          game_id TEXT, snake_id TEXT, snake_name TEXT, is_you INTEGER,
          customizations_json TEXT NOT NULL DEFAULT '{}',
          PRIMARY KEY (game_id, snake_id)
        )
      """)

      copied = copy_game_snakes(
        source, destination, batch_size=10, retained_ids={"game-1"},
      )
      rows = destination.execute("""
        SELECT snake_id, snake_name, is_you, customizations_json
        FROM game_snakes ORDER BY snake_id
      """).fetchall()
    finally:
      source.close()
      destination.close()

    self.assertEqual(copied, 2)
    self.assertEqual(rows, [
      ("snake-1", "PrismBattleSnake", 1, "{}"),
      ("snake-2", "Enemy", 0, "{}"),
    ])

  def test_copy_game_snakes_defaults_legacy_schema_to_empty_customizations(self):
    source = sqlite3.connect(":memory:")
    destination = sqlite3.connect(":memory:")
    try:
      source.execute("""
        CREATE TABLE game_snakes (
          game_id TEXT, snake_id TEXT, snake_name TEXT, is_you INTEGER
        )
      """)
      source.execute(
        "INSERT INTO game_snakes VALUES (?, ?, ?, ?)",
        ("game-1", "snake-1", "LegacySnake", 0),
      )
      source.execute("""
        CREATE TABLE snake_turns (
          game_id TEXT, snake_id TEXT, snake_name TEXT, is_you INTEGER
        )
      """)
      destination.execute("""
        CREATE TABLE game_snakes (
          game_id TEXT, snake_id TEXT, snake_name TEXT, is_you INTEGER,
          customizations_json TEXT NOT NULL DEFAULT '{}'
        )
      """)

      copied = copy_game_snakes(
        source, destination, batch_size=10, retained_ids={"game-1"},
      )
      row = destination.execute(
        "SELECT customizations_json FROM game_snakes"
      ).fetchone()
    finally:
      source.close()
      destination.close()

    self.assertEqual(copied, 1)
    self.assertEqual(row, ("{}",))

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