cleveragents-core/features/mocks/fast_init_test_helpers.py

"""Test helpers for exercising ``_fast_init_or_upgrade`` closure logic.

This module provides utilities that safely swap the
``_original_init_or_upgrade`` reference captured inside the
``_fast_init_or_upgrade`` closure (installed by
``features.environment._install_template_db_patch``) with a
``unittest.mock.MagicMock`` so that step definitions can assert
whether the original migration runner was invoked.

All helpers live in ``features/mocks/`` per the project's mock-placement
rule.
"""

from __future__ import annotations

import contextlib
import types
from collections.abc import Generator
from typing import Any
from unittest.mock import MagicMock


def _find_closure_cell(
    func: types.FunctionType,
    var_name: str,
) -> tuple[int, types.CellType] | None:
    """Locate a closure cell by free-variable name.

    Args:
        func: The function whose closure to inspect.
        var_name: The free-variable name to find.

    Returns:
        A ``(index, cell)`` tuple if found, otherwise ``None``.

    Raises:
        TypeError: If *func* is not a function with a closure.
    """
    if not isinstance(func, types.FunctionType):
        raise TypeError(f"Expected a function, got {type(func).__name__}")
    freevars: tuple[str, ...] = func.__code__.co_freevars
    closure: tuple[types.CellType, ...] | None = func.__closure__
    if closure is None:
        return None
    for idx, name in enumerate(freevars):
        if name == var_name:
            return idx, closure[idx]
    return None


@contextlib.contextmanager
def patch_original_init_or_upgrade() -> Generator[MagicMock]:
    """Context manager that replaces ``_original_init_or_upgrade`` inside
    the ``_fast_init_or_upgrade`` closure with a :class:`~unittest.mock.MagicMock`.

    Yields the mock so callers can assert on call counts and arguments.
    Restores the original cell contents on exit.

    Raises:
        RuntimeError: If the fast-init patch is not installed on
            ``MigrationRunner.init_or_upgrade`` (e.g. template DB missing).

    Example usage in a step definition::

        with patch_original_init_or_upgrade() as mock_original:
            runner.init_or_upgrade()
            assert mock_original.called
    """
    from cleveragents.infrastructure.database.migration_runner import (
        MigrationRunner,
    )

    fast_fn: Any = MigrationRunner.init_or_upgrade
    if not isinstance(fast_fn, types.FunctionType):
        raise RuntimeError(
            "MigrationRunner.init_or_upgrade is not a plain function — "
            "the _fast_init_or_upgrade patch does not appear to be installed."
        )

    result = _find_closure_cell(fast_fn, "_original_init_or_upgrade")
    if result is None:
        raise RuntimeError(
            "Cannot find '_original_init_or_upgrade' in the closure of "
            f"{fast_fn.__qualname__}. Is the template-DB patch installed?"
        )

    _idx, cell = result
    cell_ref: Any = cell
    saved: Any = cell_ref.cell_contents
    mock = MagicMock(name="_original_init_or_upgrade")
    cell_ref.cell_contents = mock
    try:
        yield mock
    finally:
        cell_ref.cell_contents = saved