cleveragents-core/features/steps/db_migration_lifecycle_steps.py

"""Step definitions for db_migration_lifecycle.feature — core lifecycle.

Exercises the full Alembic migration lifecycle: forward application,
rollback, round-trip testing, CLI wrappers, stamp logic, and the
init_database function.
"""

from __future__ import annotations

import os
import tempfile
from typing import Any
from unittest.mock import patch

from behave import given, then, when
from sqlalchemy import inspect as sa_inspect


def _create_temp_db_path() -> str:
    """Create a temporary file path for a SQLite database.

    Returns the path as a string; the caller is responsible for cleanup.
    """
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    return path


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

# Tables that are created by the Alembic migration chain.
# NOTE: ``llm_traces`` and ``audit_log`` are defined as ORM models but
# do not have dedicated migration scripts yet — they were historically
# created only via ``Base.metadata.create_all()``.  They are excluded
# from this set until their migration scripts are added.
_EXPECTED_TABLES = {
    "projects",
    "plans",
    "contexts",
    "changes",
    "debug_attempts",
    "actors",
    "actions",
    "action_invariants",
    "action_arguments",
    "v3_plans",
    "plan_projects",
    "plan_arguments",
    "plan_invariants",
    "ns_projects",
    "project_resource_links",
    "resource_types",
    "resources",
    "resource_edges",
    "resource_links",
    "tools",
    "tool_resource_bindings",
    "validation_attachments",
    "sessions",
    "session_messages",
    "automation_profiles",
    "skills",
    "skill_items",
    "locks",
    "decisions",
    "decision_dependencies",
    "checkpoint_metadata",
    "changeset_entries",
    "tool_invocations",
    "async_jobs",
    "repo_indexes",
    "indexed_files",
}


def _get_table_names(engine: Any) -> set[str]:
    inspector = sa_inspect(engine)
    return set(inspector.get_table_names())


def _make_engine(url: str = "sqlite:///:memory:") -> Any:
    from sqlalchemy import create_engine

    return create_engine(url, connect_args={"check_same_thread": False})


# ---------------------------------------------------------------------------
# Background
# ---------------------------------------------------------------------------


@given("a fresh in-memory database for migration lifecycle testing")
def step_fresh_in_memory_db(context: Any) -> None:
    context.db_url = "sqlite:///:memory:"
    context.engine = _make_engine(context.db_url)
    # Ensure the migration runner module's engine cache is clean for this test
    from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES

    MEMORY_ENGINES.pop(context.db_url, None)


# ---------------------------------------------------------------------------
# Given
# ---------------------------------------------------------------------------


@given("all migrations have been applied")
def step_all_migrations_applied(context: Any) -> None:
    from cleveragents.infrastructure.database.migration_runner import MigrationRunner

    runner = MigrationRunner(context.db_url)
    runner.run_migrations(engine=context.engine)
    context.runner = runner


@given("a database with legacy tables but no alembic_version")
def step_legacy_db_no_alembic(context: Any) -> None:
    from cleveragents.infrastructure.database.models import Base

    # Create all tables via create_all (simulates pre-Alembic state)
    Base.metadata.create_all(context.engine)
    # Verify alembic_version does NOT exist yet
    tables = _get_table_names(context.engine)
    assert "alembic_version" not in tables


# ---------------------------------------------------------------------------
# When
# ---------------------------------------------------------------------------


@when("I run all migrations forward to head")
def step_run_migrations_forward(context: Any) -> None:
    from cleveragents.infrastructure.database.migration_runner import MigrationRunner

    runner = MigrationRunner(context.db_url)
    runner.run_migrations(engine=context.engine)
    context.runner = runner


@when("I downgrade all migrations to base")
def step_downgrade_to_base(context: Any) -> None:
    from alembic import command

    runner = context.runner
    with context.engine.connect() as conn:
        runner.alembic_cfg.attributes["connection"] = conn
        try:
            command.downgrade(runner.alembic_cfg, "base")
            conn.commit()
        finally:
            runner.alembic_cfg.attributes.pop("connection", None)


@when('I downgrade to the initial migration "{revision}"')
def step_downgrade_to_revision(context: Any, revision: str) -> None:
    from alembic import command

    runner = context.runner
    with context.engine.connect() as conn:
        runner.alembic_cfg.attributes["connection"] = conn
        try:
            command.downgrade(runner.alembic_cfg, revision)
            conn.commit()
        finally:
            runner.alembic_cfg.attributes.pop("connection", None)


@when("I upgrade back to head")
def step_upgrade_to_head(context: Any) -> None:
    runner = context.runner
    runner.run_migrations(engine=context.engine)


@when("I invoke the db upgrade CLI command")
def step_invoke_db_upgrade(context: Any) -> None:

    from cleveragents.cli.main import main

    # CLI commands create their own engine, so we must use a file-based
    # DB to persist across engine instances.
    db_path = _create_temp_db_path()
    context.cli_db_path = db_path
    cli_url = f"sqlite:///{db_path}"
    with patch.dict(os.environ, {"CLEVERAGENTS_DATABASE_URL": cli_url}):
        context.cli_exit_code = main(["db", "upgrade"])
    # Re-read the DB to verify tables
    context.engine = _make_engine(cli_url)


@when("I invoke the db current CLI command")
def step_invoke_db_current(context: Any) -> None:

    from cleveragents.cli.main import main

    # Use a file-based DB so the CLI sees the migrated schema
    db_path = _create_temp_db_path()
    cli_url = f"sqlite:///{db_path}"
    # First apply migrations so there's something to report
    from cleveragents.infrastructure.database.migration_runner import MigrationRunner

    MigrationRunner(cli_url).init_or_upgrade()
    with patch.dict(os.environ, {"CLEVERAGENTS_DATABASE_URL": cli_url}):
        context.cli_exit_code = main(["db", "current"])


@when('I invoke the db downgrade CLI command with revision "{revision}"')
def step_invoke_db_downgrade(context: Any, revision: str) -> None:

    from cleveragents.cli.main import main
    from cleveragents.infrastructure.database.migration_runner import MigrationRunner

    # Use a file-based DB so the CLI sees the migrated schema
    db_path = _create_temp_db_path()
    cli_url = f"sqlite:///{db_path}"
    MigrationRunner(cli_url).init_or_upgrade()
    with patch.dict(os.environ, {"CLEVERAGENTS_DATABASE_URL": cli_url}):
        context.cli_exit_code = main(["db", "downgrade", revision])


@when("I run init_or_upgrade on the legacy database")
def step_init_or_upgrade_legacy(context: Any) -> None:
    from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES
    from cleveragents.infrastructure.database.migration_runner import MigrationRunner

    MEMORY_ENGINES[context.db_url] = context.engine
    try:
        runner = MigrationRunner(context.db_url)
        runner.init_or_upgrade()
        context.runner = runner
    finally:
        MEMORY_ENGINES.pop(context.db_url, None)


@when("I call init_database with an in-memory URL")
def step_call_init_database(context: Any) -> None:

    from cleveragents.infrastructure.database.models import init_database

    # Use a file-based temporary database so the engine can be reopened
    # and the tables inspected after init_database returns.
    db_path = _create_temp_db_path()
    file_url = f"sqlite:///{db_path}"
    context.result_engine = init_database(file_url)


# ---------------------------------------------------------------------------
# Then
# ---------------------------------------------------------------------------


@then("the database should have all expected tables")
def step_db_has_expected_tables(context: Any) -> None:
    tables = _get_table_names(context.engine)
    missing = _EXPECTED_TABLES - tables
    assert not missing, f"Missing tables: {missing}"


@then("the current revision should be the head revision")
def step_current_is_head(context: Any) -> None:
    from alembic.script import ScriptDirectory

    runner = context.runner
    head = ScriptDirectory.from_config(runner.alembic_cfg).get_current_head()

    from alembic.runtime.migration import MigrationContext

    with context.engine.connect() as conn:
        current = MigrationContext.configure(conn).get_current_revision()

    assert current == head, f"Expected {head}, got {current}"


@then("the current revision should be None")
def step_current_is_none(context: Any) -> None:
    from alembic.runtime.migration import MigrationContext

    with context.engine.connect() as conn:
        current = MigrationContext.configure(conn).get_current_revision()
    assert current is None, f"Expected None, got {current}"


@then("the database should have no application tables")
def step_db_has_no_app_tables(context: Any) -> None:
    tables = _get_table_names(context.engine)
    # Only alembic_version may remain after full downgrade
    app_tables = tables - {"alembic_version"}
    assert not app_tables, f"Unexpected tables remain: {app_tables}"


@then('the current revision should be "{expected}"')
def step_current_revision_is(context: Any, expected: str) -> None:
    from alembic.runtime.migration import MigrationContext

    with context.engine.connect() as conn:
        current = MigrationContext.configure(conn).get_current_revision()
    assert current == expected, f"Expected {expected}, got {current}"


@then("the CLI should report a successful upgrade")
def step_cli_upgrade_success(context: Any) -> None:
    assert context.cli_exit_code == 0, f"CLI exited with code {context.cli_exit_code}"


@then("the CLI should display the current revision")
def step_cli_displays_revision(context: Any) -> None:
    assert context.cli_exit_code == 0, f"CLI exited with code {context.cli_exit_code}"


@then("the CLI should report a successful downgrade")
def step_cli_downgrade_success(context: Any) -> None:
    assert context.cli_exit_code == 0, f"CLI exited with code {context.cli_exit_code}"


@then("the database should be stamped with alembic_version")
def step_db_stamped_with_alembic(context: Any) -> None:
    tables = _get_table_names(context.engine)
    assert "alembic_version" in tables, f"alembic_version not found in {tables}"


@then("the database revision should be at head")
def step_db_revision_at_head(context: Any) -> None:
    from alembic.runtime.migration import MigrationContext
    from alembic.script import ScriptDirectory

    runner = context.runner
    head = ScriptDirectory.from_config(runner.alembic_cfg).get_current_head()
    with context.engine.connect() as conn:
        current = MigrationContext.configure(conn).get_current_revision()
    assert current == head, f"Expected head {head}, got {current}"


@then("the resulting engine should have a valid schema")
def step_engine_has_valid_schema(context: Any) -> None:
    tables = _get_table_names(context.result_engine)
    # At minimum, core tables should exist
    assert "actors" in tables or "projects" in tables, f"Got tables: {tables}"


@then("the alembic_version table should exist")
def step_alembic_version_exists(context: Any) -> None:
    tables = _get_table_names(context.result_engine)
    assert "alembic_version" in tables, f"alembic_version not in {tables}"