cleveragents-core/features/mocks/test_uow_factory.py

"""Shared in-memory UnitOfWork factory for test isolation.

Provides ``build_test_uow`` — a factory that constructs an in-memory
SQLite-backed ``UnitOfWork`` suitable for both Behave step definitions
and Robot Framework integration test helpers.

Centralising the UoW construction eliminates duplication and ensures
both test suites exercise the same database setup.  Any changes to
``UnitOfWork.__init__`` must be reflected here in a single place rather
than in each consuming test file.

Behave Context Adapters:
- ``use_test_uow(context)`` — Attach a test UoW to a Behave context and
  register engine cleanup on scenario teardown.

Used by:
- ``features/steps/plan_use_action_args_integrity_steps.py``
- ``features/steps/plan_service_steps.py``
- ``features/steps/project_cli_steps.py``
- ``features/steps/project_service_coverage_boost_steps.py``
- ``robot/helper_plan_use_action_args_integrity.py``
"""

from __future__ import annotations

import contextlib
from collections.abc import Callable
from typing import Any

from sqlalchemy import create_engine, event
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

from cleveragents.infrastructure.database.models import Base
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork


def build_test_uow() -> tuple[UnitOfWork, sessionmaker[Session], Engine]:
    """Build an in-memory UoW for testing.

    Creates a fresh SQLite ``:memory:`` database with all tables, foreign
    key enforcement enabled, and a ``UnitOfWork`` instance that bypasses
    the constructor to avoid running migrations inside the test harness.

    Returns:
        A three-tuple of ``(uow, session_factory, engine)``.
    """
    engine = create_engine(
        "sqlite:///:memory:",
        echo=False,
        future=True,
        connect_args={"check_same_thread": False},
    )

    @event.listens_for(engine, "connect")
    def _enable_fk(dbapi_conn: Any, _rec: Any) -> None:
        cursor = dbapi_conn.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

    Base.metadata.create_all(engine)
    sf: sessionmaker[Session] = sessionmaker(
        bind=engine,
        expire_on_commit=False,
        autoflush=False,
        autocommit=False,
        class_=Session,
    )

    uow = UnitOfWork.__new__(UnitOfWork)
    # Keep these attributes in sync with UnitOfWork.__init__ — we bypass
    # the constructor to avoid running migrations inside the test harness.
    uow.database_url = "sqlite:///:memory:"
    uow._engine = engine
    uow._session_factory = sf
    uow._database_initialized = True
    uow._prompt_for_migration = None
    uow._require_confirmation = False

    return uow, sf, engine


def use_test_uow(
    context: Any,
    seed_hook: Callable[[UnitOfWork], None] | None = None,
) -> UnitOfWork:
    """Attach a test UnitOfWork to a Behave context with automatic cleanup.

    This is the recommended way to set up database fixtures in Behave step
    definitions. It ensures:
    - Consistent in-memory SQLite configuration across all tests
    - Foreign key PRAGMA enforcement
    - Automatic engine cleanup on scenario teardown
    - Optional seed hook for test data setup

    Args:
        context: The Behave context object (typically from a step function).
        seed_hook: Optional callable that receives the UoW and populates
                   test data. Called immediately after UoW creation.

    Returns:
        The configured UnitOfWork instance, also attached to context.uow.

    Example:
        @given("a plan service with test database")
        def step_plan_service_with_test_db(context):
            uow = use_test_uow(context)
            context.plan_service = PlanService(uow)
    """
    uow, sf, engine = build_test_uow()

    # Attach to context for use in step definitions
    context.uow = uow
    context.session_factory = sf
    context.engine = engine

    # Register cleanup callback
    def cleanup_engine() -> None:
        """Dispose of the engine connection pool."""
        engine.dispose()

    # Store cleanup function for later invocation
    if not hasattr(context, "_cleanup_callbacks"):
        context._cleanup_callbacks = []
    context._cleanup_callbacks.append(cleanup_engine)

    # Run optional seed hook
    if seed_hook is not None:
        seed_hook(uow)

    return uow


def cleanup_test_uow(context: Any) -> None:
    """Clean up test UnitOfWork resources attached to a Behave context.

    This should be called in an ``after_scenario`` hook or similar teardown.

    Args:
        context: The Behave context object.
    """
    if hasattr(context, "_cleanup_callbacks"):
        for callback in context._cleanup_callbacks:
            with contextlib.suppress(Exception):
                callback()
        context._cleanup_callbacks.clear()

    # Clear references
    if hasattr(context, "uow"):
        delattr(context, "uow")
    if hasattr(context, "session_factory"):
        delattr(context, "session_factory")
    if hasattr(context, "engine"):
        delattr(context, "engine")