"""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}"