turnstone/tests/test_state_writer.py

"""Unit tests for ``turnstone.core.state_writer``.

Tests cover the contract callers depend on:

* Buffered transitions coalesce per ws_id (last state wins).
* ``flush_now=True`` bypasses the buffer (used for terminal ERROR
  transitions and any other write that must be durable on return).
* ``discard`` drops pending and waits for any in-flight flush to
  complete (the bug-3 invariant — close()'s sync ``closed`` write must
  not be overtaken by a buffered transient).
* Bounded buffer evicts oldest under capacity pressure.
* DB error during flush doesn't poison the loop; subsequent flushes
  still run.
* Shutdown drains any pending entries synchronously.
"""

from __future__ import annotations

import threading
import time

from turnstone.core.state_writer import StateWriter


class _FakeStorage:
    """Records update_workstream_state calls. Optionally raises or pauses."""

    def __init__(self, *, raises: BaseException | None = None) -> None:
        self.calls: list[tuple[str, str]] = []
        self.raises = raises
        self._call_lock = threading.Lock()
        # Optional gate to pin a write inside update_workstream_state
        # so the test can race ``discard`` against an in-flight flush.
        self.write_gate: threading.Event | None = None
        # Set by the writer thread once it enters update_workstream_state.
        self.write_started = threading.Event()

    def update_workstream_state(self, ws_id: str, state: str) -> None:
        if self.write_gate is not None:
            self.write_started.set()
            self.write_gate.wait(timeout=2.0)
        with self._call_lock:
            self.calls.append((ws_id, state))
        if self.raises is not None:
            raise self.raises


def _drain(writer: StateWriter) -> None:
    """Trigger a single flush synchronously."""
    writer.flush()


# ---------------------------------------------------------------------------
# Coalescing + flush
# ---------------------------------------------------------------------------


def test_buffered_transitions_coalesce_per_ws_id() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage)

    writer.record("ws-1", "thinking")
    writer.record("ws-1", "running")
    writer.record("ws-1", "idle")
    writer.record("ws-2", "thinking")

    _drain(writer)
    # Only the latest state per ws_id should land.
    assert sorted(storage.calls) == sorted([("ws-1", "idle"), ("ws-2", "thinking")])


def test_flush_now_bypasses_buffer_and_writes_sync() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage)

    # Pre-buffer something for a different ws_id to prove the sync
    # path doesn't drain the whole buffer.
    writer.record("ws-other", "running")

    writer.record("ws-err", "error", flush_now=True)
    # ws-err landed sync, ws-other still buffered.
    assert ("ws-err", "error") in storage.calls
    assert ("ws-other", "running") not in storage.calls

    _drain(writer)
    assert ("ws-other", "running") in storage.calls


def test_flush_now_swallows_storage_error() -> None:
    storage = _FakeStorage(raises=RuntimeError("db down"))
    writer = StateWriter(storage)
    # Should not raise — set_state path can't recover from a storage
    # write failure mid-transition.
    writer.record("ws-1", "error", flush_now=True)


def test_flush_now_drops_pending_buffered_state_for_same_ws_id() -> None:
    """Terminal-bypass invariant: a buffered transient for the same
    ws_id must NOT flush AFTER the sync ``flush_now`` write and
    clobber the terminal state. (This was a real correctness gap
    flagged by /review.)"""
    storage = _FakeStorage()
    writer = StateWriter(storage)

    # Buffer a transient transition first.
    writer.record("ws-A", "running")

    # Sync ERROR write must drop the buffered 'running' AND wait on
    # the flush_lock so any in-flight flush can't sneak through after.
    writer.record("ws-A", "error", flush_now=True)

    # Run the flusher; nothing pending for ws-A any more.
    writer.flush()

    ws_writes = [s for w, s in storage.calls if w == "ws-A"]
    # The sync 'error' must be in storage, and 'running' must NOT have
    # been flushed AFTER it.
    assert "error" in ws_writes, ws_writes
    assert ws_writes[-1] == "error", f"buffered 'running' clobbered terminal 'error': {ws_writes}"
    # Stronger: the 'running' should never have landed at all.
    assert "running" not in ws_writes, ws_writes


def test_flush_now_waits_for_in_flight_flush_to_complete() -> None:
    """Same shape as the discard wait: if a flusher is mid-write on
    the same ws_id, ``flush_now`` must NOT issue its sync write
    until the flusher finishes — otherwise the order on the wire is
    flush_now → flusher's late write → final state is the transient,
    not the terminal."""
    storage = _FakeStorage()
    storage.write_gate = threading.Event()
    writer = StateWriter(storage)

    writer.record("ws-A", "running")

    flush_done = threading.Event()

    def _flush_in_bg() -> None:
        writer.flush()
        flush_done.set()

    flusher = threading.Thread(target=_flush_in_bg, daemon=True)
    flusher.start()
    assert storage.write_started.wait(timeout=1.0)

    flush_now_done = threading.Event()

    def _flush_now_in_bg() -> None:
        writer.record("ws-A", "error", flush_now=True)
        flush_now_done.set()

    fn_thread = threading.Thread(target=_flush_now_in_bg, daemon=True)
    fn_thread.start()
    time.sleep(0.05)
    assert flush_now_done.is_set() is False, (
        "flush_now returned before in-flight flush released flush_lock"
    )

    storage.write_gate.set()
    flusher.join(timeout=2.0)
    fn_thread.join(timeout=2.0)
    assert flush_done.is_set() and flush_now_done.is_set()
    # The flusher's 'running' lands first, then flush_now's 'error'.
    ws_writes = [s for w, s in storage.calls if w == "ws-A"]
    assert ws_writes == ["running", "error"], ws_writes


# ---------------------------------------------------------------------------
# Bounded buffer
# ---------------------------------------------------------------------------


def test_bounded_buffer_evicts_oldest_on_capacity() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage, max_buffer=3)

    writer.record("ws-1", "running")
    writer.record("ws-2", "running")
    writer.record("ws-3", "running")
    # ws-4 forces eviction of ws-1 (oldest).
    writer.record("ws-4", "running")

    _drain(writer)
    landed = {ws_id for ws_id, _ in storage.calls}
    assert "ws-1" not in landed
    assert {"ws-2", "ws-3", "ws-4"} <= landed


def test_bounded_buffer_update_existing_does_not_evict() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage, max_buffer=2)

    writer.record("ws-1", "running")
    writer.record("ws-2", "running")
    # Update existing — must not evict.
    writer.record("ws-1", "idle")

    _drain(writer)
    landed = dict(storage.calls)
    assert landed["ws-1"] == "idle"
    assert landed["ws-2"] == "running"


# ---------------------------------------------------------------------------
# Resilience
# ---------------------------------------------------------------------------


def test_storage_error_does_not_poison_subsequent_flushes() -> None:
    storage = _FakeStorage(raises=RuntimeError("db blip"))
    errors: list[Exception] = []
    writer = StateWriter(storage, on_flush_error=errors.append)

    writer.record("ws-1", "running")
    _drain(writer)
    # Error was surfaced via callback.
    assert len(errors) == 1

    # Storage recovers; next flush succeeds.
    storage.raises = None
    writer.record("ws-2", "idle")
    _drain(writer)
    assert ("ws-2", "idle") in storage.calls


# ---------------------------------------------------------------------------
# discard / close-race
# ---------------------------------------------------------------------------


def test_discard_drops_pending_buffered_state() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage)

    writer.record("ws-close", "running")
    writer.discard("ws-close")
    _drain(writer)
    assert storage.calls == []


def test_discard_waits_for_in_flight_flush_to_complete() -> None:
    """The bug-3 invariant: ``close()`` calls ``discard`` BEFORE its
    sync ``state='closed'`` write. If a flusher was mid-write for the
    same ws_id, the flusher's write must complete BEFORE
    ``discard`` returns — so ``close()``'s sync write strictly
    follows the flusher's transient write, leaving 'closed' as the
    final state. (If discard returned early, close's 'closed' write
    could be overwritten by the flusher's late 'running' write.)
    """
    storage = _FakeStorage()
    storage.write_gate = threading.Event()
    writer = StateWriter(storage)

    writer.record("ws-A", "running")

    # Kick off a flush in a background thread; it will block inside
    # update_workstream_state on storage.write_gate.
    flush_done = threading.Event()

    def _flush_in_bg() -> None:
        writer.flush()
        flush_done.set()

    flusher = threading.Thread(target=_flush_in_bg, daemon=True)
    flusher.start()
    assert storage.write_started.wait(timeout=1.0)
    assert flush_done.is_set() is False  # writer is pinned

    # Call discard concurrently — it must NOT return until the flush
    # completes.
    discard_done = threading.Event()

    def _discard_in_bg() -> None:
        writer.discard("ws-A")
        discard_done.set()

    discarder = threading.Thread(target=_discard_in_bg, daemon=True)
    discarder.start()
    # discard should be blocked on flush_lock.
    time.sleep(0.05)
    assert discard_done.is_set() is False, "discard returned before flusher released the write"

    # Release the writer; both threads should complete now.
    storage.write_gate.set()
    flusher.join(timeout=2.0)
    discarder.join(timeout=2.0)
    assert flush_done.is_set()
    assert discard_done.is_set()
    # The flusher's write went through.
    assert ("ws-A", "running") in storage.calls


# ---------------------------------------------------------------------------
# Lifecycle
# ---------------------------------------------------------------------------


def test_start_starts_flusher_and_buffered_writes_land() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage, flush_interval=0.05)
    writer.start()
    try:
        writer.record("ws-1", "running")
        # Wait up to 1s for the flusher to drain.
        for _ in range(20):
            if storage.calls:
                break
            time.sleep(0.05)
        assert ("ws-1", "running") in storage.calls
    finally:
        writer.shutdown(timeout=2.0)


def test_shutdown_drains_pending_synchronously() -> None:
    storage = _FakeStorage()
    # Long flush interval so no automatic drain happens.
    writer = StateWriter(storage, flush_interval=60.0)
    writer.start()
    try:
        writer.record("ws-1", "running")
        writer.record("ws-2", "thinking")
    finally:
        writer.shutdown(timeout=2.0)
    landed = {ws_id for ws_id, _ in storage.calls}
    assert {"ws-1", "ws-2"} <= landed


def test_start_is_idempotent() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage, flush_interval=0.05)
    writer.start()
    first_thread = writer._thread
    writer.start()
    assert writer._thread is first_thread
    writer.shutdown(timeout=2.0)


def test_discard_times_out_when_flush_hangs() -> None:
    """If the flusher is wedged on a stuck Postgres connection, discard
    must NOT block forever — callers hold ws._lock across this call,
    so an unbounded wait would deadlock all close paths system-wide."""
    storage = _FakeStorage()
    storage.write_gate = threading.Event()  # never released
    writer = StateWriter(storage)

    writer.record("ws-A", "running")

    # Pin the flusher inside update_workstream_state.
    flusher = threading.Thread(target=writer.flush, daemon=True)
    flusher.start()
    assert storage.write_started.wait(timeout=1.0)

    # discard must return within ~timeout, NOT hang forever.
    start = time.monotonic()
    writer.discard("ws-A", flush_lock_timeout=0.1)
    elapsed = time.monotonic() - start
    assert elapsed < 1.0, f"discard hung: {elapsed:.2f}s"

    # Cleanup.
    storage.write_gate.set()
    flusher.join(timeout=2.0)


def test_shutdown_is_idempotent() -> None:
    storage = _FakeStorage()
    writer = StateWriter(storage, flush_interval=0.05)
    writer.start()
    writer.shutdown(timeout=2.0)
    # Second shutdown is a no-op, must not raise.
    writer.shutdown(timeout=2.0)


# ---------------------------------------------------------------------------
# Wake-on-record
# ---------------------------------------------------------------------------


def test_record_wakes_flusher_immediately() -> None:
    """Single transitions get persisted within ~one round-trip rather
    than waiting up to flush_interval seconds."""
    storage = _FakeStorage()
    # Long interval — only the wake event should drive the flush.
    writer = StateWriter(storage, flush_interval=10.0)
    writer.start()
    try:
        writer.record("ws-1", "running")
        for _ in range(30):
            if storage.calls:
                break
            time.sleep(0.02)
        assert ("ws-1", "running") in storage.calls
    finally:
        writer.shutdown(timeout=2.0)