chat/tests/test_rewind.py

"""Tests for Task 28 — rewind with snapshot, impact preview, and re-projection.

Per Requirements §10.1, rewind must:

* take a pre-rewind snapshot of all projected tables (so the user can recover),
* truncate the event log past a chosen event id,
* clear projected tables and re-project from the truncated log so live state
  matches "what the world looked like at turn N" (no stale rows from rewound
  events).

These tests cover the functional core. The HTTP route surface is left to the
plan's polish pass — tests exercise via direct service calls.
"""

from __future__ import annotations

import json

from chat.db.connection import open_db
from chat.db.migrate import apply_migrations
from chat.eventlog.log import append_event
from chat.eventlog.projector import project
from chat.services.rewind import compute_rewind_preview, execute_rewind
from chat.services.snapshot import take_snapshot

# Importing the state modules registers their projector handlers as a
# side effect — the test would otherwise see an unprojected db after
# re-projection because the registry would be empty.
import chat.state.entities  # noqa: F401
import chat.state.edges  # noqa: F401
import chat.state.memory  # noqa: F401
import chat.state.world  # noqa: F401
import chat.state.manual_edit  # noqa: F401


def _seed_5_turns(db):
    """Seed: bot + chat + 5 mock user/assistant turn pairs.

    user_turn / assistant_turn have no projector handlers — they live in
    the event_log purely for transcript rendering — so the only
    projection-bearing events are bot_authored and chat_created. That
    makes the post-rewind invariants easy to assert without needing the
    real classifier pass.
    """
    apply_migrations(db)
    with open_db(db) as conn:
        append_event(
            conn,
            kind="bot_authored",
            payload={
                "id": "bot_a",
                "name": "BotA",
                "persona": "...",
                "voice_samples": [],
                "traits": [],
                "backstory": "",
                "initial_relationship_to_you": "",
                "kickoff_prose": "",
            },
        )
        append_event(
            conn,
            kind="chat_created",
            payload={
                "id": "chat_bot_a",
                "host_bot_id": "bot_a",
                "initial_time": "2026-04-26T20:00:00+00:00",
                "narrative_anchor": "Day 1",
                "weather": "",
            },
        )
        for i in range(5):
            append_event(
                conn,
                kind="user_turn",
                payload={
                    "chat_id": "chat_bot_a",
                    "prose": f"turn {i}",
                    "segments": [],
                },
            )
            append_event(
                conn,
                kind="assistant_turn",
                payload={
                    "chat_id": "chat_bot_a",
                    "speaker_id": "bot_a",
                    "text": f"reply {i}",
                    "truncated": False,
                    "user_turn_id": i,
                },
            )
        project(conn)


def test_take_snapshot_writes_file_to_disk(tmp_path):
    db = tmp_path / "t.db"
    _seed_5_turns(db)
    with open_db(db) as conn:
        snapshot_path = take_snapshot(
            conn, data_dir=tmp_path / "data", kind="rewind"
        )
    assert snapshot_path.exists()
    assert snapshot_path.parent == tmp_path / "data" / "snapshots" / "rewind"
    data = json.loads(snapshot_path.read_text())
    assert "event_log" in data
    assert "bots" in data
    assert "chats" in data
    # Bot is in the dump
    assert any(b["id"] == "bot_a" for b in data["bots"])


def test_compute_rewind_preview_counts_kinds(tmp_path):
    db = tmp_path / "t.db"
    _seed_5_turns(db)
    with open_db(db) as conn:
        # The 1st assistant_turn is the 4th event:
        # 1 bot_authored, 2 chat_created, 3 user_turn, 4 assistant_turn.
        # Everything past it should be in the preview (4 user + 4 assistant = 8).
        first_assistant = conn.execute(
            "SELECT id FROM event_log WHERE kind='assistant_turn' "
            "ORDER BY id LIMIT 1"
        ).fetchone()[0]
        preview = compute_rewind_preview(conn, after_event_id=first_assistant)
        assert preview["after_event_id"] == first_assistant
        assert preview["total_events"] > 0
        # by_kind sums to total_events.
        assert sum(preview["by_kind"].values()) == preview["total_events"]
        # We rewound past assistant_turn #1, so 4 user + 4 assistant remain.
        assert preview["by_kind"].get("user_turn") == 4
        assert preview["by_kind"].get("assistant_turn") == 4


def test_execute_rewind_truncates_and_reprojects(tmp_path):
    db = tmp_path / "t.db"
    data_dir = tmp_path / "data"
    _seed_5_turns(db)
    with open_db(db) as conn:
        first_assistant = conn.execute(
            "SELECT id FROM event_log WHERE kind='assistant_turn' "
            "ORDER BY id LIMIT 1"
        ).fetchone()[0]

    snapshot_path = execute_rewind(
        db_path=db, data_dir=data_dir, after_event_id=first_assistant
    )
    # Snapshot is written under data/snapshots/rewind/.
    assert snapshot_path.exists()
    assert snapshot_path.parent == data_dir / "snapshots" / "rewind"

    # Verify event_log truncated and projected state matches state-at-turn.
    with open_db(db) as conn:
        max_id = conn.execute("SELECT MAX(id) FROM event_log").fetchone()[0]
        assert max_id == first_assistant
        # Bot still exists (re-projected from preserved bot_authored event).
        bot = conn.execute(
            "SELECT id FROM bots WHERE id = 'bot_a'"
        ).fetchone()
        assert bot is not None
        # Chat still exists (re-projected from preserved chat_created event).
        chat = conn.execute(
            "SELECT id FROM chats WHERE id = 'chat_bot_a'"
        ).fetchone()
        assert chat is not None