344 lines
11 KiB
Python
344 lines
11 KiB
Python
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("DROP TABLE scheduled_prompts")
|
|
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
|
|
applied: list[int] = []
|
|
# Use versions above current LATEST_VERSION so none have been applied yet.
|
|
next_version = migrations.LATEST_VERSION + 1
|
|
monkeypatch.setattr(
|
|
migrations,
|
|
"MIGRATIONS",
|
|
[
|
|
*migrations.MIGRATIONS,
|
|
(next_version, lambda conn: applied.append(next_version)),
|
|
(next_version + 1, lambda conn: applied.append(next_version + 1)),
|
|
],
|
|
)
|
|
monkeypatch.setattr(migrations, "LATEST_VERSION", next_version + 1)
|
|
engine = get_engine(settings)
|
|
migrate(engine)
|
|
assert applied == [next_version, next_version + 1]
|
|
assert _user_version(engine) == next_version + 1
|
|
migrate(engine) # re-run is a no-op
|
|
assert applied == [next_version, next_version + 1]
|
|
|
|
|
|
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 _scheduled_prompt_column_names(engine) -> list[str]:
|
|
with engine.connect() as conn:
|
|
rows = conn.exec_driver_sql("PRAGMA table_info(scheduled_prompts)").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 "sandbox_json" in cols
|
|
assert _user_version(engine) == migrations.LATEST_VERSION
|
|
|
|
|
|
def test_v1_db_upgraded_to_latest(tmp_path):
|
|
# Build a v1 database: create all tables, drop v2+ objects, 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("DROP TABLE scheduled_prompts")
|
|
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 "scheduled_prompts" in inspect(engine).get_table_names()
|
|
assert _user_version(engine) == migrations.LATEST_VERSION
|
|
|
|
|
|
def test_fresh_db_has_scheduled_prompts_table(tmp_path):
|
|
settings = _tmp_settings(tmp_path)
|
|
ensure_db(settings)
|
|
engine = get_engine(settings)
|
|
|
|
tables = inspect(engine).get_table_names()
|
|
assert "scheduled_prompts" in tables
|
|
cols = _scheduled_prompt_column_names(engine)
|
|
assert cols == [
|
|
"id",
|
|
"session_id",
|
|
"prompt",
|
|
"due_at",
|
|
"status",
|
|
"error",
|
|
"created_at",
|
|
"sent_at",
|
|
]
|
|
assert _user_version(engine) == migrations.LATEST_VERSION
|
|
|
|
|
|
def test_v2_db_upgraded_to_latest_adds_scheduled_prompts_and_sandbox(tmp_path):
|
|
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("DROP TABLE scheduled_prompts")
|
|
conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN sandbox_json")
|
|
conn.exec_driver_sql("PRAGMA user_version = 2")
|
|
|
|
migrate(engine)
|
|
|
|
assert "scheduled_prompts" in inspect(engine).get_table_names()
|
|
assert "sandbox_json" in _session_column_names(engine)
|
|
assert _user_version(engine) == migrations.LATEST_VERSION
|
|
|
|
|
|
def test_scheduled_prompt_one_row_per_session(tmp_path):
|
|
from datetime import datetime, timedelta
|
|
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from hqt.db.models import ScheduledPrompt
|
|
|
|
settings = _tmp_settings(tmp_path)
|
|
ensure_db(settings)
|
|
factory = get_session_factory(get_engine(settings))
|
|
|
|
with factory() as session:
|
|
p = Project(name="sched-proj", path="/tmp/sched-proj")
|
|
h = Harness(name="claude-code-scheduled")
|
|
session.add_all([p, h])
|
|
session.flush()
|
|
s = Session(
|
|
project_id=p.id,
|
|
harness_id=h.id,
|
|
tmux_session_name="hqt-scheduled",
|
|
)
|
|
session.add(s)
|
|
session.flush()
|
|
due = datetime.now() + timedelta(minutes=30)
|
|
session.add(ScheduledPrompt(session_id=s.id, prompt="continue", due_at=due))
|
|
session.commit()
|
|
|
|
session.add(ScheduledPrompt(session_id=s.id, prompt="again", due_at=due))
|
|
with pytest.raises(IntegrityError):
|
|
session.commit()
|
|
|
|
|
|
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_four():
|
|
assert migrations.LATEST_VERSION == 4
|
|
|
|
|
|
def test_migration_adds_sandbox_json_column(tmp_path):
|
|
# Simulate a v3 DB whose sessions table predates the sandbox 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 = 3")
|
|
migrate(engine)
|
|
cols = [c["name"] for c in inspect(engine).get_columns("sessions")]
|
|
assert "sandbox_json" in cols
|
|
assert _user_version(engine) == migrations.LATEST_VERSION
|
|
|
|
|
|
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}'
|