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 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 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 == "001_initial_schema" 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