5.5Testing Migrations in Python

import subprocess
import psycopg2
import pytest

@pytest.fixture
def migration_db():
    """Create a disposable database for migration testing."""
    db_name = "migration_test_" + str(int(time.time()))

    # Create database from production schema
    subprocess.run(["createdb", db_name], check=True)
    subprocess.run(
        f"pg_dump --schema-only production_db | psql {db_name}",
        shell=True, check=True
    )

    conn = psycopg2.connect(dbname=db_name)
    yield conn, db_name

    conn.close()
    subprocess.run(["dropdb", db_name], check=True)

def test_migration_applies_cleanly(migration_db):
    conn, db_name = migration_db
    # Apply migration
    subprocess.run(
        f"psql {db_name} < migrations/0042_add_phone_column.sql",
        shell=True, check=True
    )
    # Verify new column exists
    cursor = conn.cursor()
    cursor.execute("""
        SELECT column_name, data_type, is_nullable
        FROM information_schema.columns
        WHERE table_name = 'users' AND column_name = 'phone'
    """)
    row = cursor.fetchone()
    assert row is not None, "Phone column was not created"
    assert row[1] == "character varying"  # Expected data type

def test_migration_preserves_data(migration_db):
    conn, db_name = migration_db
    cursor = conn.cursor()

    # Insert known data before migration
    cursor.execute("""
        INSERT INTO users (name, email) VALUES
        ('Alice', 'alice@test.com'),
        ('Bob', 'bob@test.com')
    """)
    conn.commit()

    # Apply migration
    subprocess.run(
        f"psql {db_name} < migrations/0042_add_phone_column.sql",
        shell=True, check=True
    )

    # Verify data is preserved
    conn = psycopg2.connect(dbname=db_name)  # Reconnect
    cursor = conn.cursor()
    cursor.execute("SELECT name, email FROM users ORDER BY name")
    rows = cursor.fetchall()
    assert len(rows) == 2
    assert rows[0][0] == "Alice"
    assert rows[1][0] == "Bob"

def test_migration_rollback(migration_db):
    conn, db_name = migration_db

    # Apply migration
    subprocess.run(
        f"psql {db_name} < migrations/0042_add_phone_column.sql",
        shell=True, check=True
    )

    # Rollback
    subprocess.run(
        f"psql {db_name} < migrations/0042_rollback.sql",
        shell=True, check=True
    )

    # Verify column is removed
    conn = psycopg2.connect(dbname=db_name)
    cursor = conn.cursor()
    cursor.execute("""
        SELECT column_name FROM information_schema.columns
        WHERE table_name = 'users' AND column_name = 'phone'
    """)
    assert cursor.fetchone() is None, "Phone column should not exist after rollback"

def test_not_null_column_has_default(migration_db):
    conn, db_name = migration_db
    cursor = conn.cursor()

    # Insert data before migration
    cursor.execute("INSERT INTO users (name, email) VALUES ('Test', 'test@test.com')")
    conn.commit()

    # Apply migration that adds NOT NULL column
    subprocess.run(
        f"psql {db_name} < migrations/0043_add_status_column.sql",
        shell=True, check=True
    )

    # Verify existing rows have the default value
    conn = psycopg2.connect(dbname=db_name)
    cursor = conn.cursor()
    cursor.execute("SELECT status FROM users WHERE email = 'test@test.com'")
    status = cursor.fetchone()[0]
    assert status is not None, "Existing rows should have default status value"
    assert status == "active", "Default status should be 'active'"