mvcc-isolation-tests

# tests/db/test_isolation.py
import os, threading, time
import psycopg
import pytest
from psycopg import IsolationLevel

DSN = os.environ.get("TEST_DSN", "dbname=test")
WAIT = 5  # seconds; a wait that times out fails the test instead of hanging

SCHEMA = """
DROP TABLE IF EXISTS accounts, orders, doctors;
CREATE TABLE accounts (id text PRIMARY KEY, balance int NOT NULL);
CREATE TABLE orders (id serial PRIMARY KEY, status text NOT NULL);
CREATE TABLE doctors (name text PRIMARY KEY, status text NOT NULL);
INSERT INTO accounts VALUES ('acc1', 100);
INSERT INTO orders (status) VALUES ('pending'), ('pending'), ('pending');
INSERT INTO doctors VALUES ('alice', 'on_call'), ('bob', 'on_call');
"""

def scalar(sql, *args):
    with psycopg.connect(DSN, autocommit=True) as c:
        return c.execute(sql, args).fetchone()[0]

def on_call():
    return scalar("SELECT count(*) FROM doctors WHERE status = 'on_call'")

@pytest.fixture(autouse=True)
def fresh_tables():
    """Rebuild the tables before every test and refuse to run on unexpected state."""
    with psycopg.connect(DSN, autocommit=True) as c:
        c.execute(SCHEMA)
    assert scalar("SELECT balance FROM accounts WHERE id = 'acc1'") == 100
    assert scalar("SELECT count(*) FROM orders") == 3
    assert on_call() == 2

def run_pair(work_a, work_b, level):
    """Run two workloads on two connections at `level`; return (results, errors)."""
    conns = [psycopg.connect(DSN) for _ in range(2)]
    for c in conns:
        c.isolation_level = level
    results, errors = {}, {}

    def run(key, conn, work):
        try:
            results[key] = work(conn)
        except Exception as e:          # recorded, asserted by the test
            errors[key] = e
            conn.rollback()

    threads = [threading.Thread(target=run, args=("a", conns[0], work_a)),
               threading.Thread(target=run, args=("b", conns[1], work_b))]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=3 * WAIT)
    assert not any(t.is_alive() for t in threads), "workload hung (lock wait?)"
    for c in conns:
        c.close()
    return results, errors

def test_non_repeatable_read_under_read_committed():
    read1, wrote = threading.Event(), threading.Event()

    def reader(conn):
        first = conn.execute("SELECT balance FROM accounts WHERE id = 'acc1'").fetchone()[0]
        read1.set()
        assert wrote.wait(WAIT)
        second = conn.execute("SELECT balance FROM accounts WHERE id = 'acc1'").fetchone()[0]
        conn.commit()
        return first, second

    def writer(conn):
        assert read1.wait(WAIT)
        conn.execute("UPDATE accounts SET balance = 200 WHERE id = 'acc1'")
        conn.commit()
        wrote.set()

    results, errors = run_pair(reader, writer, IsolationLevel.READ_COMMITTED)
    assert not errors
    assert results["a"] == (100, 200)   # Read Committed allows it

def test_no_phantom_under_repeatable_read():
    read1, wrote = threading.Event(), threading.Event()

    def reader(conn):
        q = "SELECT count(*) FROM orders WHERE status = 'pending'"
        first = conn.execute(q).fetchone()[0]
        read1.set()
        assert wrote.wait(WAIT)
        second = conn.execute(q).fetchone()[0]
        conn.commit()
        return first, second

    def writer(conn):
        assert read1.wait(WAIT)
        conn.execute("INSERT INTO orders (status) VALUES ('pending')")
        conn.commit()
        wrote.set()

    results, errors = run_pair(reader, writer, IsolationLevel.REPEATABLE_READ)
    assert not errors
    assert results["a"] == (3, 3)       # PostgreSQL's snapshot hides the insert
    assert scalar("SELECT count(*) FROM orders") == 4

def take_off(name, both_read):
    def work(conn):
        count = conn.execute("SELECT count(*) FROM doctors WHERE status = 'on_call'").fetchone()[0]
        both_read.wait()                # both transactions have read before either writes
        if count >= 2:
            conn.execute("UPDATE doctors SET status = 'off' WHERE name = %s", (name,))
        conn.commit()
    return work

def test_write_skew_happens_under_repeatable_read():
    both_read = threading.Barrier(2, timeout=WAIT)
    _, errors = run_pair(take_off("alice", both_read), take_off("bob", both_read),
                         IsolationLevel.REPEATABLE_READ)
    assert not errors
    assert on_call() == 0               # the anomaly: nobody on call

def test_serializable_aborts_one_of_the_pair():
    both_read = threading.Barrier(2, timeout=WAIT)
    _, errors = run_pair(take_off("alice", both_read), take_off("bob", both_read),
                         IsolationLevel.SERIALIZABLE)
    assert len(errors) == 1
    assert isinstance(next(iter(errors.values())), psycopg.errors.SerializationFailure)
    assert on_call() == 1

def with_serializable_retry(conn, fn, max_attempts=5):
    """Run fn(conn, attempt) in a SERIALIZABLE transaction, retrying serialization failures."""
    conn.isolation_level = IsolationLevel.SERIALIZABLE
    for attempt in range(1, max_attempts + 1):
        try:
            with conn.transaction():
                return fn(conn, attempt)
        except psycopg.errors.SerializationFailure:
            if attempt == max_attempts:
                raise
            time.sleep(0.01 * 2 ** attempt)

def test_retry_completes_and_keeps_the_invariant():
    both_read = threading.Barrier(2, timeout=WAIT)
    attempts = []

    def take_off_retrying(name):
        def fn(conn, attempt):
            attempts.append((name, attempt))
            count = conn.execute("SELECT count(*) FROM doctors WHERE status = 'on_call'").fetchone()[0]
            if attempt == 1:
                both_read.wait()        # force the conflict only on the first try
            if count >= 2:
                conn.execute("UPDATE doctors SET status = 'off' WHERE name = %s", (name,))
        return lambda conn: with_serializable_retry(conn, fn)

    _, errors = run_pair(take_off_retrying("alice"), take_off_retrying("bob"),
                         IsolationLevel.SERIALIZABLE)
    assert not errors                   # both calls completed
    assert len(attempts) == 3           # exactly one retry
    assert on_call() == 1

TEST_DSN="host=localhost dbname=test_isolation user=postgres" pytest tests/db/test_isolation.py