hqt/tests/test_db.py

from pathlib import Path

import pytest
from sqlalchemy import inspect

from hqt.config import Settings
from hqt.db import migrations
from hqt.db.engine import ensure_db, get_engine, get_session_factory
from hqt.db.migrations import migrate
from hqt.db.models import Base, Harness, Project, Session


def _tmp_settings(tmp_path: Path) -> Settings:
    return Settings(db_path=tmp_path / "test.db")


def test_ensure_db_creates_tables(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    engine = get_engine(settings)
    tables = inspect(engine).get_table_names()
    assert "projects" in tables
    assert "sessions" in tables
    assert "harnesses" in tables
    assert "mcp_servers" in tables
    assert "project_mcp_servers" in tables
    assert "project_skills" in tables


def test_project_crud(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    engine = get_engine(settings)
    factory = get_session_factory(engine)

    with factory() as session:
        p = Project(name="test", path="/tmp/test")
        session.add(p)
        session.commit()
        assert p.id is not None

        p.name = "updated"
        session.commit()

        fetched = session.get(Project, p.id)
        assert fetched.name == "updated"

        session.delete(fetched)
        session.commit()
        assert session.get(Project, p.id) is None


def test_session_with_fk(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    engine = get_engine(settings)
    factory = get_session_factory(engine)

    with factory() as session:
        p = Project(name="proj", path="/tmp/proj")
        h = Harness(name="claude-code")
        session.add_all([p, h])
        session.commit()

        s = Session(
            project_id=p.id,
            harness_id=h.id,
            tmux_session_name="hqt-test-1",
        )
        session.add(s)
        session.commit()

        assert s.project.name == "proj"
        assert s.harness.name == "claude-code"
        assert s in p.sessions


def _user_version(engine) -> int:
    with engine.connect() as conn:
        return conn.exec_driver_sql("PRAGMA user_version").scalar() or 0


def test_fresh_db_stamped_latest_version(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    assert _user_version(get_engine(settings)) == migrations.LATEST_VERSION


def test_unversioned_existing_db_stamped_latest(tmp_path):
    # Simulates a user DB created before versioning existed: tables present,
    # user_version still 0.  migrate() should apply all pending migrations and
    # stamp the DB at LATEST_VERSION (not just BASELINE_VERSION).
    settings = _tmp_settings(tmp_path)
    settings.db_path.parent.mkdir(parents=True, exist_ok=True)
    engine = get_engine(settings)
    Base.metadata.create_all(engine)
    with engine.begin() as conn:
        conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_path")
        conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_branch")
        conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN sandbox_json")
    migrate(engine)
    assert _user_version(engine) == migrations.LATEST_VERSION


def test_db_from_newer_hqt_raises(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    engine = get_engine(settings)
    with engine.begin() as conn:
        conn.exec_driver_sql("PRAGMA user_version = 9999")
    with pytest.raises(RuntimeError, match="newer"):
        migrate(engine)


def test_pending_migrations_applied_in_order_once(tmp_path, monkeypatch):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)  # stamps migrations.LATEST_VERSION (currently 3)
    applied: list[int] = []
    # Use versions above current LATEST_VERSION so none have been applied yet.
    monkeypatch.setattr(
        migrations,
        "MIGRATIONS",
        [
            *migrations.MIGRATIONS,
            (4, lambda conn: applied.append(4)),
            (5, lambda conn: applied.append(5)),
        ],
    )
    monkeypatch.setattr(migrations, "LATEST_VERSION", 5)
    engine = get_engine(settings)
    migrate(engine)
    assert applied == [4, 5]
    assert _user_version(engine) == 5
    migrate(engine)  # re-run is a no-op
    assert applied == [4, 5]


def _session_column_names(engine) -> list[str]:
    with engine.connect() as conn:
        rows = conn.exec_driver_sql("PRAGMA table_info(sessions)").fetchall()
    return [row[1] for row in rows]


def test_fresh_db_has_worktree_columns(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    engine = get_engine(settings)
    cols = _session_column_names(engine)
    assert "worktree_path" in cols
    assert "worktree_branch" in cols
    assert _user_version(engine) == 3


def test_v1_db_upgraded_to_v2(tmp_path):
    # Build a v1 database: create all tables, drop v2+ columns, stamp v1.
    settings = _tmp_settings(tmp_path)
    settings.db_path.parent.mkdir(parents=True, exist_ok=True)
    engine = get_engine(settings)
    Base.metadata.create_all(engine)
    with engine.begin() as conn:
        conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_path")
        conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_branch")
        conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN sandbox_json")
        conn.exec_driver_sql("PRAGMA user_version = 1")

    migrate(engine)

    cols = _session_column_names(engine)
    assert "worktree_path" in cols
    assert "worktree_branch" in cols
    assert "sandbox_json" in cols
    assert _user_version(engine) == 3


def test_session_worktree_fields_default_none(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    engine = get_engine(settings)
    factory = get_session_factory(engine)

    with factory() as session:
        p = Project(name="wt-proj", path="/tmp/wt-proj")
        h = Harness(name="claude-code-wt")
        session.add_all([p, h])
        session.commit()

        # Session without worktree fields
        s_plain = Session(
            project_id=p.id,
            harness_id=h.id,
            tmux_session_name="hqt-plain",
        )
        # Session with worktree fields
        s_wt = Session(
            project_id=p.id,
            harness_id=h.id,
            tmux_session_name="hqt-worktree",
            worktree_path="/tmp/wt",
            worktree_branch="feature/x",
        )
        session.add_all([s_plain, s_wt])
        session.commit()

        assert s_plain.worktree_path is None
        assert s_plain.worktree_branch is None
        assert s_wt.worktree_path == "/tmp/wt"
        assert s_wt.worktree_branch == "feature/x"


def test_rows_stay_readable_after_session_closes(tmp_path):
    # Services use per-operation sessions; rows they return must remain
    # readable after the originating session closes (expire_on_commit=False).
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    factory = get_session_factory(get_engine(settings))
    with factory() as session:
        p = Project(name="x", path="/x")
        session.add(p)
        session.commit()
    assert p.name == "x"  # would raise DetachedInstanceError without the flag


def test_latest_version_is_three():
    assert migrations.LATEST_VERSION == 3


def test_migration_adds_sandbox_json_column(tmp_path):
    # Simulate a baseline DB whose sessions table predates the column.
    settings = _tmp_settings(tmp_path)
    settings.db_path.parent.mkdir(parents=True, exist_ok=True)
    engine = get_engine(settings)
    with engine.begin() as conn:
        conn.exec_driver_sql(
            "CREATE TABLE sessions ("
            "id INTEGER PRIMARY KEY, project_id INTEGER, harness_id INTEGER, "
            "tmux_session_name TEXT)"
        )
        conn.exec_driver_sql("PRAGMA user_version = 2")
    migrate(engine)
    cols = [c["name"] for c in inspect(engine).get_columns("sessions")]
    assert "sandbox_json" in cols
    assert _user_version(engine) == 3


def test_session_round_trips_sandbox_json(tmp_path):
    settings = _tmp_settings(tmp_path)
    ensure_db(settings)
    factory = get_session_factory(get_engine(settings))
    with factory() as session:
        p = Project(name="proj", path="/tmp/proj-sandbox")
        h = Harness(name="claude-sandbox")
        session.add_all([p, h])
        session.commit()
        s = Session(
            project_id=p.id,
            harness_id=h.id,
            tmux_session_name="hqt-9",
            sandbox_json='{"fs": "rw", "net": true}',
        )
        session.add(s)
        session.commit()
        assert session.get(Session, s.id).sandbox_json == '{"fs": "rw", "net": true}'