forked from cleveragents/cleveragents-core
457 lines
16 KiB
Python
457 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from behave import given, then, when
|
|
|
|
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
|