cleveragents-core/features/steps/migration_runner_steps.py

from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from typing import Any
from unittest.mock import MagicMock, patch

from behave import given, then, when

from cleveragents.core.exceptions import MigrationNotApprovedError
from cleveragents.infrastructure.database.migration_runner import (
    MEMORY_ENGINES,
    MigrationRunner,
)


class FakeConnection:
    def __init__(self) -> None:
        self.entered = False
        self.exit_called = False
        self.closed_direct = False
        self.committed = False

    def __enter__(self) -> FakeConnection:
        self.entered = True
        return self

    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
        self.exit_called = True
        return False

    def close(self) -> None:
        self.closed_direct = True

    def commit(self) -> None:
        self.committed = True


class FakeEngine:
    def __init__(self) -> None:
        self.connections: list[FakeConnection] = []
        self.disposed = False

    def connect(self) -> FakeConnection:
        conn = FakeConnection()
        self.connections.append(conn)
        return conn

    def dispose(self) -> None:
        self.disposed = True


@dataclass
class LegacyStampContext:
    stamp_calls: list[Any]
    connection_flags: list[bool]


@given('a migration runner configured for "{database_url}"')
def step_given_migration_runner(context, database_url: str) -> None:
    context.database_url = database_url
    context.runner = MigrationRunner(database_url)


@when("I attempt to load the alembic config without an alembic.ini file")
def step_when_load_missing_alembic(context) -> None:
    context.runner._alembic_cfg = None

    from pathlib import Path

    def fake_exists(self) -> bool:
        if str(self).endswith("alembic.ini"):
            return False
        return Path.exists(self)

    context.alembic_error = None
    try:
        with patch.object(Path, "exists", fake_exists):
            _ = context.runner.alembic_cfg
    except FileNotFoundError as exc:
        context.alembic_error = exc


@then('a FileNotFoundError should be raised mentioning "{expected_snippet}"')
def step_then_missing_alembic_error(context, expected_snippet: str) -> None:
    assert context.alembic_error is not None, "Expected FileNotFoundError"
    assert expected_snippet in str(context.alembic_error)


@when("I run migrations without providing an engine")
def step_when_run_migrations_no_engine(context) -> None:
    env_var = "CLEVERAGENTS_DATABASE_URL"
    original_env = os.environ.get(env_var)
    os.environ[env_var] = "sqlite:///pre-existing.db"
    context.env_before = os.environ[env_var]
    captured_connection_attr: list[bool] = []

    def fake_upgrade(cfg, revision) -> None:
        captured_connection_attr.append("connection" in cfg.attributes)

    with patch(
        "cleveragents.infrastructure.database.migration_runner.command.upgrade",
        side_effect=fake_upgrade,
    ):
        context.runner.run_migrations()

    context.upgrade_call_count = len(captured_connection_attr)
    context.connection_attribute_present = (
        captured_connection_attr[0] if captured_connection_attr else None
    )
    context.connection_attribute_state = dict(context.runner.alembic_cfg.attributes)
    context.env_after = os.environ.get(env_var)

    if original_env is None:
        os.environ.pop(env_var, None)
    else:
        os.environ[env_var] = original_env


@then("the upgrade command should be invoked without a connection attribute")
def step_then_upgrade_without_connection(context) -> None:
    assert context.upgrade_call_count == 1
    assert context.connection_attribute_present is False
    assert "connection" not in context.connection_attribute_state


@then("the CLEVERAGENTS_DATABASE_URL environment variable should be restored")
def step_then_env_restored(context) -> None:
    assert context.env_after == context.env_before


@when("I initialize or upgrade the database with cached in-memory engine")
def step_when_init_in_memory(context) -> None:
    memory_snapshot = dict(MEMORY_ENGINES)
    MEMORY_ENGINES.clear()

    fake_engine = FakeEngine()
    inspector = MagicMock()
    inspector.get_table_names.return_value = []
    create_calls: list[Any] = []

    def fake_create_engine(url: str, **kwargs: Any) -> FakeEngine:
        create_calls.append((url, kwargs))
        return fake_engine

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            side_effect=fake_create_engine,
        ),
        patch("sqlalchemy.inspect", return_value=inspector),
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.upgrade"
        ) as upgrade_mock,
    ):
        context.runner.init_or_upgrade()
        context.runner.init_or_upgrade()
        upgrade_calls = list(upgrade_mock.call_args_list)

    context.upgrade_call_count_with_engine = len(upgrade_calls)
    context.create_engine_call_count = len(create_calls)
    context.fake_engine = fake_engine
    context.fake_engine_connections = list(fake_engine.connections)
    context.connection_attrs_after_init = dict(context.runner.alembic_cfg.attributes)
    context.memory_cache_state = dict(MEMORY_ENGINES)

    MEMORY_ENGINES.clear()
    MEMORY_ENGINES.update(memory_snapshot)


@then("migrations should run using the cached engine connection")
def step_then_cached_engine_used(context) -> None:
    assert context.upgrade_call_count_with_engine == 2
    assert context.create_engine_call_count == 1
    assert len(context.fake_engine_connections) == 4
    even_connections = context.fake_engine_connections[::2]
    odd_connections = context.fake_engine_connections[1::2]
    for conn in even_connections:
        assert conn.exit_called is True
    for conn in odd_connections:
        assert conn.closed_direct is True
    assert "connection" not in context.connection_attrs_after_init
    cache_entry = context.memory_cache_state.get(context.runner.database_url)
    assert cache_entry is context.fake_engine


@then("the cached engine should remain available without disposal")
def step_then_cached_engine_not_disposed(context) -> None:
    assert context.fake_engine.disposed is False


@when("I initialize or upgrade the database with legacy tables present")
def step_when_legacy_stamping(context) -> None:
    fake_engine = FakeEngine()
    inspector = MagicMock()
    inspector.get_table_names.return_value = ["users"]
    stamp_context = LegacyStampContext(stamp_calls=[], connection_flags=[])

    def fake_create_engine(url: str, **kwargs: Any) -> FakeEngine:
        context.legacy_create_engine_call = (url, kwargs)
        return fake_engine

    def fake_stamp(cfg, revision) -> None:
        stamp_context.stamp_calls.append((cfg, revision))
        stamp_context.connection_flags.append("connection" in cfg.attributes)

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            side_effect=fake_create_engine,
        ),
        patch("sqlalchemy.inspect", return_value=inspector),
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.stamp",
            side_effect=fake_stamp,
        ) as stamp_mock,
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.upgrade"
        ) as upgrade_mock,
    ):
        context.runner.init_or_upgrade()
        context.legacy_stamp_calls = list(stamp_mock.call_args_list)
        context.legacy_upgrade_calls = list(upgrade_mock.call_args_list)

    context.legacy_stamp_connection_has_attr = any(stamp_context.connection_flags)
    context.legacy_stamp_revision = (
        stamp_context.stamp_calls[0][1] if stamp_context.stamp_calls else None
    )
    context.legacy_fake_engine = fake_engine
    context.legacy_connection_attr_after = dict(context.runner.alembic_cfg.attributes)


@then("the stamp command should run using the active connection")
def step_then_stamp_uses_connection(context) -> None:
    assert len(context.legacy_stamp_calls) == 1
    assert context.legacy_stamp_revision == "head"
    assert context.legacy_stamp_connection_has_attr is True
    assert len(context.legacy_upgrade_calls) == 0
    assert "connection" not in context.legacy_connection_attr_after


@then("the external database engine should be disposed after initialization")
def step_then_external_engine_disposed(context) -> None:
    assert context.legacy_fake_engine.disposed is True


@when("I request pending migrations for a database with no current revision")
def step_when_pending_no_current(context) -> None:
    revisions = ["rev_003", "rev_002", None]

    class FakeRevision:
        def __init__(self, revision: Any) -> None:
            self.revision = revision

    fake_revisions = [FakeRevision(r) for r in revisions]

    def walk_revisions():
        return iter(fake_revisions)

    fake_script_dir = MagicMock()
    fake_script_dir.walk_revisions.side_effect = walk_revisions

    with (
        patch.object(MigrationRunner, "get_current_revision", return_value=None),
        patch(
            "cleveragents.infrastructure.database.migration_runner.ScriptDirectory.from_config",
            return_value=fake_script_dir,
        ),
    ):
        context.pending_migrations = context.runner.get_pending_migrations()


@then("the pending migration list should be ordered from oldest to newest")
def step_then_pending_migrations_ordered(context) -> None:
    assert context.pending_migrations == [None, "rev_002", "rev_003"]


@when("I request the current revision from the database")
def step_when_get_current_revision(context) -> None:
    fake_engine = FakeEngine()
    migration_context = MagicMock()
    migration_context.get_current_revision.return_value = "001_initial_schema"

    def fake_create_engine(url: str, **kwargs: Any) -> FakeEngine:
        context.current_rev_create_call = (url, kwargs)
        return fake_engine

    def fake_configure(conn):
        context.current_rev_connection = conn
        return migration_context

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            side_effect=fake_create_engine,
        ),
        patch(
            "cleveragents.infrastructure.database.migration_runner.MigrationContext.configure",
            side_effect=fake_configure,
        ),
    ):
        context.current_revision = context.runner.get_current_revision()

    context.current_rev_fake_engine = fake_engine


@then("the migration context should be queried for the current revision")
def step_then_migration_context_queried(context) -> None:
    assert context.current_revision == "001_initial_schema"


@then("the temporary connection should be closed afterward")
def step_then_temp_connection_closed(context) -> None:
    assert len(context.current_rev_fake_engine.connections) == 1
    assert context.current_rev_fake_engine.connections[0].exit_called is True


@when("I initialize or upgrade a file-based SQLite database")
def step_when_init_file_based_sqlite(context) -> None:
    import shutil
    from pathlib import Path

    # Extract db path from URL
    db_path = context.database_url.replace("sqlite:///", "")
    if not db_path.startswith("/"):
        db_path = "/" + db_path
    db_file = Path(db_path)

    # Clean up if exists
    if db_file.parent.exists():
        shutil.rmtree(db_file.parent)

    context.db_parent_dir = db_file.parent

    fake_engine = FakeEngine()
    inspector = MagicMock()
    inspector.get_table_names.return_value = []
    context.file_sqlite_create_calls = []

    def fake_create_engine(url: str, **kwargs: Any) -> FakeEngine:
        context.file_sqlite_create_calls.append((url, kwargs))
        return fake_engine

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            side_effect=fake_create_engine,
        ),
        patch("sqlalchemy.inspect", return_value=inspector),
        patch("cleveragents.infrastructure.database.migration_runner.command.upgrade"),
    ):
        context.runner.init_or_upgrade()

    context.file_sqlite_fake_engine = fake_engine


@then("the parent directory should be created if it does not exist")
def step_then_parent_dir_created(context) -> None:
    assert context.db_parent_dir.exists()


@then("the database engine should be disposed after initialization")
def step_then_engine_disposed_after_init(context) -> None:
    assert context.file_sqlite_fake_engine.disposed is True


@then("the engine should be created with check_same_thread set to False")
def step_then_engine_created_with_args(context) -> None:
    assert len(context.file_sqlite_create_calls) > 0
    _url, kwargs = context.file_sqlite_create_calls[0]
    assert "connect_args" in kwargs
    assert kwargs["connect_args"]["check_same_thread"] is False


@when("I initialize the database and migrations are already applied")
def step_when_init_with_existing_migrations(context) -> None:
    fake_engine = FakeEngine()
    inspector = MagicMock()
    inspector.get_table_names.return_value = ["alembic_version", "users"]
    context.up_to_date_upgrade_calls = []

    def fake_create_engine(url: str, **kwargs: Any) -> FakeEngine:
        return fake_engine

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            side_effect=fake_create_engine,
        ),
        patch("sqlalchemy.inspect", return_value=inspector),
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.upgrade"
        ) as upgrade_mock,
        patch.object(MigrationRunner, "get_pending_migrations", return_value=[]),
    ):
        context.check_migrations_result = context.runner.check_migrations_needed()
        context.runner.init_or_upgrade()
        context.up_to_date_upgrade_calls = list(upgrade_mock.call_args_list)

    context.up_to_date_fake_engine = fake_engine


@then("no additional migrations should be run")
def step_then_no_migrations_run(context) -> None:
    assert len(context.up_to_date_upgrade_calls) == 0


@then("check migrations needed should return False")
def step_then_check_migrations_false(context) -> None:
    assert context.check_migrations_result is False


@when("I initialize the database with pending migrations detected")
def step_when_pending_migrations_detected(context) -> None:
    memory_snapshot = dict(MEMORY_ENGINES)
    fake_engine = FakeEngine()
    MEMORY_ENGINES[context.runner.database_url] = fake_engine

    inspector = MagicMock()
    inspector.get_table_names.return_value = ["alembic_version", "users"]

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            return_value=fake_engine,
        ) as create_engine_mock,
        patch("sqlalchemy.inspect", return_value=inspector),
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.upgrade"
        ) as upgrade_mock,
        patch.object(
            MigrationRunner, "get_pending_migrations", return_value=["rev_002"]
        ),
    ):
        context.runner.init_or_upgrade()
        context.pending_upgrade_calls = list(upgrade_mock.call_args_list)
        context.pending_create_engine_calls = create_engine_mock.call_count

    context.pending_fake_engine = fake_engine
    context.pending_engine_connections = list(fake_engine.connections)
    context.pending_memory_cache_state = dict(MEMORY_ENGINES)

    MEMORY_ENGINES.clear()
    MEMORY_ENGINES.update(memory_snapshot)


@then("run migrations should be invoked with the existing engine")
def step_then_run_migrations_existing_engine(context) -> None:
    assert len(context.pending_upgrade_calls) == 1
    assert context.pending_create_engine_calls == 0
    assert len(context.pending_engine_connections) == 2
    first_conn, second_conn = context.pending_engine_connections
    assert first_conn.exit_called is True
    assert second_conn.closed_direct is True
    cache_entry = context.pending_memory_cache_state.get(context.runner.database_url)
    assert cache_entry is context.pending_fake_engine


@then("the in-memory engine should not be disposed")
def step_then_pending_engine_not_disposed(context) -> None:
    assert context.pending_fake_engine.disposed is False


@when("I initialize the database with pending migrations requiring confirmation")
def step_when_pending_migrations_require_confirmation(context) -> None:
    inspector = MagicMock()
    inspector.get_table_names.return_value = ["alembic_version", "users"]
    prompt_calls: list[str] = []

    def prompt(message: str) -> bool:
        prompt_calls.append(message)
        return True

    fake_engine = FakeEngine()

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            return_value=fake_engine,
        ),
        patch("sqlalchemy.inspect", return_value=inspector),
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.upgrade"
        ) as upgrade_mock,
        patch.object(
            MigrationRunner,
            "get_pending_migrations",
            return_value=["rev_002", "rev_003"],
        ),
    ):
        context.runner.init_or_upgrade(
            require_confirmation=True, prompt_for_migration=prompt
        )
        context.migration_prompt_calls = list(prompt_calls)
        context.migration_upgrade_call_count = upgrade_mock.call_count


@then("the migration approval prompt should be invoked")
def step_then_prompt_invoked(context) -> None:
    assert len(getattr(context, "migration_prompt_calls", [])) == 1


@then("migrations should run after approval")
def step_then_migrations_run_after_approval(context) -> None:
    assert getattr(context, "migration_upgrade_call_count", 0) == 1


@when("I decline migration approval when pending migrations exist")
def step_when_decline_migration(context) -> None:
    inspector = MagicMock()
    inspector.get_table_names.return_value = ["alembic_version"]

    def prompt(_message: str) -> bool:
        return False

    fake_engine = FakeEngine()

    with (
        patch(
            "cleveragents.infrastructure.database.migration_runner.create_engine",
            return_value=fake_engine,
        ),
        patch("sqlalchemy.inspect", return_value=inspector),
        patch(
            "cleveragents.infrastructure.database.migration_runner.command.upgrade"
        ) as upgrade_mock,
        patch.object(
            MigrationRunner, "get_pending_migrations", return_value=["rev_002"]
        ),
    ):
        context.migration_error = None
        try:
            context.runner.init_or_upgrade(
                require_confirmation=True, prompt_for_migration=prompt
            )
        except Exception as exc:
            context.migration_error = exc
        context.migration_decline_upgrade_calls = upgrade_mock.call_count


@then("a migration approval error should be raised")
def step_then_migration_error(context) -> None:
    assert isinstance(context.migration_error, MigrationNotApprovedError)
    assert context.migration_decline_upgrade_calls == 0


def _snapshot_env(keys: list[str]) -> dict[str, str | None]:
    return {key: os.environ.get(key) for key in keys}


def _restore_env(snapshot: dict[str, str | None]) -> None:
    for key, value in snapshot.items():
        if value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = value


@when("I evaluate the default migration prompt while CI is set")
def step_when_prompt_ci(context) -> None:
    snapshot = _snapshot_env(
        ["CI", "BEHAVE_TESTING", "CLEVERAGENTS_AUTO_APPLY_MIGRATIONS"]
    )
    try:
        os.environ["CI"] = "true"
        os.environ.pop("BEHAVE_TESTING", None)
        os.environ.pop("CLEVERAGENTS_AUTO_APPLY_MIGRATIONS", None)
        context.prompt_result = MigrationRunner._default_prompt_for_migration(
            "Apply migrations in CI?"
        )
        context.prompt_confirm_calls = []
    finally:
        _restore_env(snapshot)


@then("the default prompt should auto approve without interaction")
def step_then_prompt_auto_approve_ci(context) -> None:
    assert context.prompt_result is True
    assert getattr(context, "prompt_confirm_calls", []) == []


@when("I evaluate the default migration prompt in an interactive TTY")
def step_when_prompt_tty(context) -> None:
    snapshot = _snapshot_env(
        ["CI", "BEHAVE_TESTING", "CLEVERAGENTS_AUTO_APPLY_MIGRATIONS"]
    )
    fake_stdin = MagicMock()
    fake_stdin.isatty.return_value = True
    confirm_calls: list[Any] = []

    def confirm(*args: Any, **kwargs: Any) -> bool:
        confirm_calls.append((args, kwargs))
        return False

    fake_typer = MagicMock()
    fake_typer.confirm.side_effect = confirm
    try:
        os.environ.pop("CI", None)
        os.environ.pop("BEHAVE_TESTING", None)
        os.environ["CLEVERAGENTS_AUTO_APPLY_MIGRATIONS"] = ""
        with (
            patch("sys.stdin", fake_stdin),
            patch.dict(sys.modules, {"typer": fake_typer}),
        ):
            context.prompt_result = MigrationRunner._default_prompt_for_migration(
                "Apply migrations interactively?"
            )
    finally:
        _restore_env(snapshot)

    context.prompt_confirm_calls = list(confirm_calls)


@then("the default prompt should return the typer confirmation result")
def step_then_prompt_returns_typer(context) -> None:
    assert context.prompt_result is False
    assert len(getattr(context, "prompt_confirm_calls", [])) == 1
    args, kwargs = context.prompt_confirm_calls[0]
    assert kwargs.get("default") is False
    assert "Apply migrations interactively?" in args[0]


@when("the default migration prompt raises an error while prompting")
def step_when_prompt_errors(context) -> None:
    snapshot = _snapshot_env(
        ["CI", "BEHAVE_TESTING", "CLEVERAGENTS_AUTO_APPLY_MIGRATIONS"]
    )
    fake_stdin = MagicMock()
    fake_stdin.isatty.return_value = True
    confirm_calls: list[Any] = []

    def confirm(*args: Any, **kwargs: Any) -> bool:
        confirm_calls.append((args, kwargs))
        raise RuntimeError("prompt failure")

    error_typer = MagicMock()
    error_typer.confirm.side_effect = confirm
    try:
        os.environ.pop("CI", None)
        os.environ.pop("BEHAVE_TESTING", None)
        os.environ["CLEVERAGENTS_AUTO_APPLY_MIGRATIONS"] = ""
        with (
            patch("sys.stdin", fake_stdin),
            patch.dict(sys.modules, {"typer": error_typer}),
        ):
            context.prompt_result = MigrationRunner._default_prompt_for_migration(
                "Apply migrations after errors?"
            )
    finally:
        _restore_env(snapshot)

    context.prompt_error_confirm_calls = list(confirm_calls)


@then("the default prompt should auto approve after the failure")
def step_then_prompt_auto_after_error(context) -> None:
    assert context.prompt_result is True
    assert len(getattr(context, "prompt_error_confirm_calls", [])) == 1


@given("the MigrationRunner class is available")
def step_given_migration_runner_available(context) -> None:
    """Make MigrationRunner available in context for testing alembic.ini location."""
    context.migration_runner_class = MigrationRunner


@when("I call _find_alembic_ini() to locate the configuration file")
def step_when_find_alembic_ini(context) -> None:
    """Call the _find_alembic_ini method to locate alembic.ini from package."""
    try:
        context.alembic_ini_path = MigrationRunner._find_alembic_ini()
        context.find_alembic_error = None
    except FileNotFoundError as exc:
        context.find_alembic_error = exc
        context.alembic_ini_path = None


@then("it should find alembic.ini at the package location")
def step_then_alembic_ini_found_at_package(context) -> None:
    """Verify alembic.ini is found in the package location."""
    assert context.find_alembic_error is None, (
        f"Expected to find alembic.ini but got error: {context.find_alembic_error}"
    )
    assert context.alembic_ini_path is not None
    assert context.alembic_ini_path.exists(), (
        f"alembic.ini path {context.alembic_ini_path} does not exist"
    )
    assert "migrations" in str(context.alembic_ini_path).lower(), (
        f"Expected alembic.ini in migrations directory, got {context.alembic_ini_path}"
    )


@then("the alembic.ini should be readable and valid")
def step_then_alembic_ini_readable(context) -> None:
    """Verify alembic.ini is readable and contains expected configuration."""
    assert context.alembic_ini_path.is_file()
    content = context.alembic_ini_path.read_text(encoding="utf-8")
    assert "[loggers]" in content, "alembic.ini missing loggers section"
    assert "script_location" in content, "alembic.ini missing script_location"