cleveragents-core/features/environment.py

"""Behave environment setup for feature tests."""

import contextlib
import os
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Any

LANGSMITH_ENV_VARS = [
    "CLEVERAGENTS_LANGSMITH_ENABLED",
    "CLEVERAGENTS_LANGSMITH_PROJECT",
    "CLEVERAGENTS_LANGSMITH_ENDPOINT",
    "CLEVERAGENTS_LANGSMITH_API_KEY",
    "CLEVERAGENTS_LANGSMITH_TRACING_V2",
    "CLEVERAGENTS_LANGSMITH_USER_ID",
    "LANGCHAIN_TRACING_V2",
    "LANGCHAIN_PROJECT",
    "LANGCHAIN_ENDPOINT",
    "LANGCHAIN_API_KEY",
    "LANGSMITH_PROJECT",
    "LANGSMITH_ENDPOINT",
    "LANGSMITH_API_KEY",
    "LANGSMITH_TRACING_V2",
    "LANGSMITH_USER_ID",
]


def before_all(context):
    """Set up test environment before all tests."""
    # Add src to path
    sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
    # Ensure tests never block on migration prompts or real providers
    os.environ.setdefault("CLEVERAGENTS_AUTO_APPLY_MIGRATIONS", "true")
    os.environ.setdefault("CLEVERAGENTS_TESTING_USE_MOCK_AI", "true")
    # Use per-process unique database paths so parallel test subprocesses
    # (behave-parallel) never contend on the same SQLite file.
    if "CLEVERAGENTS_DATABASE_URL" not in os.environ:
        os.environ["CLEVERAGENTS_DATABASE_URL"] = (
            f"sqlite:///{tempfile.mktemp(suffix='.db', prefix='cleveragents_')}"
        )
    if "CLEVERAGENTS_TEST_DATABASE_URL" not in os.environ:
        os.environ["CLEVERAGENTS_TEST_DATABASE_URL"] = (
            f"sqlite:///{tempfile.mktemp(suffix='.db', prefix='cleveragents_test_')}"
        )
    os.environ.setdefault("BEHAVE_TESTING", "true")

    # Set up mock AI provider for all tests
    try:
        from cleveragents.application.container import override_providers
        from features.mocks.mock_ai_provider import MockAIProvider

        # Override the AI provider with mock for all tests
        mock_provider = MockAIProvider()
        override_providers(ai_provider=mock_provider)
    except ImportError:
        pass  # Container not needed for all tests

    # --- Eliminate retry waits ---
    # Tenacity retry decorators (database_retry, network_retry, etc.) use
    # real time.sleep() waits during retries.  In tests, mocked operations
    # fail deterministically so waiting is pure overhead.  Patch
    # time.sleep() globally so all tenacity waits (and any other sleeps)
    # complete instantly.  The small handful of sleep() calls that exist
    # in step definitions already use sub-100ms waits and are unaffected
    # by this optimisation in practice.
    _install_fast_sleep_patch()

    # --- Template-DB fast-path ---
    # When CLEVERAGENTS_TEMPLATE_DB is set (by nox sessions), monkey-patch
    # MigrationRunner.init_or_upgrade so that fresh file-based SQLite
    # databases are created by copying the pre-migrated template (~1ms)
    # instead of running 25 Alembic migrations (~0.5-3s each).
    #
    # If running outside of nox (i.e. direct `behave` invocation), auto-
    # create the template DB so the fast-path is always active.
    _ensure_template_db()
    _install_template_db_patch()


def _install_fast_sleep_patch() -> None:
    """Cap ``time.sleep`` and ``asyncio.sleep`` at 10 ms for fast test execution.

    Tenacity retry decorators (``@database_retry``, ``@retry_network_operation``,
    etc.) ultimately call ``time.sleep()`` with waits of 0.5-30 s between retry
    attempts.  Async retry helpers (``retry_auto_debug``,
    ``async_retry_with_exponential_backoff``) call ``asyncio.sleep()`` with
    exponential waits of 1-4 s per attempt.  In the test suite, mocked
    operations fail deterministically, so the long sleeps are pure overhead
    (~1 s per retry cycle x hundreds of scenarios = minutes of wasted time).

    Both functions are replaced with capped versions (≤ 10 ms).  The originals
    are saved as ``time._original_sleep`` / ``asyncio._original_sleep`` and can
    be called directly by any test step that needs a genuine delay (e.g.
    CircuitBreaker recovery-timeout tests that need real wall-clock advancement
    past a 100 ms threshold).
    """
    import asyncio
    import time

    _MAX_SLEEP = 0.01  # 10 ms cap

    # --- synchronous time.sleep ---
    if not callable(getattr(time, "_original_sleep", None)):
        time._original_sleep = time.sleep  # type: ignore[attr-defined]

        def _capped_sleep(seconds: float) -> None:
            time._original_sleep(min(seconds, _MAX_SLEEP))  # type: ignore[attr-defined]

        time.sleep = _capped_sleep  # type: ignore[assignment]

    # --- asynchronous asyncio.sleep ---
    if not callable(getattr(asyncio, "_original_sleep", None)):
        asyncio._original_sleep = asyncio.sleep  # type: ignore[attr-defined]

        async def _capped_async_sleep(seconds: float, result: object = None) -> object:
            return await asyncio._original_sleep(min(seconds, _MAX_SLEEP), result)  # type: ignore[attr-defined]

        asyncio.sleep = _capped_async_sleep  # type: ignore[assignment]


def _ensure_template_db() -> None:
    """Auto-create the template DB when CLEVERAGENTS_TEMPLATE_DB is not set.

    When running tests directly via ``behave`` (without nox), the env var
    is missing.  This function creates the template on the fly using
    ``scripts/create_template_db.py`` so the fast-path is always active.
    """
    if os.environ.get("CLEVERAGENTS_TEMPLATE_DB"):
        return  # Already set by nox or CI

    template_path = Path(__file__).parent.parent / "build" / ".template-migrated.db"
    if template_path.is_file():
        # Template already exists from a prior run — reuse it.
        os.environ["CLEVERAGENTS_TEMPLATE_DB"] = str(template_path.resolve())
        return

    try:
        # Import the template creation script
        scripts_dir = Path(__file__).parent.parent / "scripts"
        sys.path.insert(0, str(scripts_dir))
        from create_template_db import create_template

        create_template(str(template_path))
        os.environ["CLEVERAGENTS_TEMPLATE_DB"] = str(template_path.resolve())
    except Exception:
        pass  # Fall back to normal Alembic migrations


def _install_template_db_patch() -> None:
    """Monkey-patch MigrationRunner to skip Alembic migrations in tests.

    For file-based SQLite: copies a pre-migrated template DB (~1 ms).
    For in-memory SQLite: uses ``Base.metadata.create_all()`` (~5 ms) instead
    of running 25 sequential Alembic migrations (~0.5-3 s).
    """
    template_path = os.environ.get("CLEVERAGENTS_TEMPLATE_DB")
    if not template_path or not Path(template_path).is_file():
        return

    try:
        from cleveragents.infrastructure.database.migration_runner import (
            MigrationRunner,
        )
    except ImportError:
        return

    _original_init_or_upgrade = MigrationRunner.init_or_upgrade

    # Prefixes used by before_scenario and step files when creating temp DBs.
    # "cleveragents_" / "cleveragents_test_" — before_scenario databases
    # "test_" — databases created inside step files (services_coverage, etc.)
    _SCENARIO_DB_PREFIXES = ("cleveragents_", "cleveragents_test_", "test_", "db.")

    def _fast_init_or_upgrade(self: Any, **kwargs: Any) -> None:
        """Replace Alembic migrations with fast alternatives.

        - Non-SQLite databases → fall through to original
        - In-memory SQLite → ``Base.metadata.create_all()`` + alembic stamp
        - File-based SQLite with matching prefix → copy template
        - Everything else → fall through to original
        """
        db_url: str = getattr(self, "database_url", "")

        # Non-SQLite: always fall through
        if not db_url.startswith("sqlite"):
            return _original_init_or_upgrade(self, **kwargs)

        # In-memory SQLite: fall through — these are rare in tests and the
        # engine hasn't been created yet at this point (UnitOfWork is lazy).
        if ":memory:" in db_url or db_url == "sqlite://":
            return _original_init_or_upgrade(self, **kwargs)

        # Extract the file path from the URL
        db_file_path = db_url.replace("sqlite:///", "")
        if not db_file_path.startswith("/"):
            db_file_path = "/" + db_file_path

        db_path = Path(db_file_path)

        # Only apply to scenario-generated temp DBs (avoid hijacking
        # migration-runner unit tests that use custom URLs).
        if not any(db_path.name.startswith(p) for p in _SCENARIO_DB_PREFIXES):
            return _original_init_or_upgrade(self, **kwargs)

        # Only copy template for databases that don't exist yet or are empty
        # (SQLite auto-creates a 0-byte file on first engine open).
        if db_path.exists() and db_path.stat().st_size > 0:
            return _original_init_or_upgrade(self, **kwargs)

        # Copy the template — creates a fully-migrated DB in ~1ms
        db_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(template_path, db_file_path)

    MigrationRunner.init_or_upgrade = _fast_init_or_upgrade  # type: ignore[assignment]


def before_scenario(context, scenario):
    """Set up before each scenario."""
    # Store original working directory
    context.original_cwd = os.getcwd()

    # Initialize cleanup list
    context._cleanup_handlers = []
    context.stubbed_clients = {}

    # Reset error attributes to prevent stale state leaking between scenarios.
    context.acms_error = None
    context.assemble_error = None

    # Ensure mock AI flag is always set so plan service tests can resolve actors
    os.environ["CLEVERAGENTS_TESTING_USE_MOCK_AI"] = "true"

    # Re-establish protective env vars every scenario.  These are set once
    # in before_all, but individual scenarios (e.g. migration_runner tests)
    # temporarily remove them.  If any cleanup path fails to restore them
    # the remaining scenarios in a serial run would be affected.
    os.environ["BEHAVE_TESTING"] = "true"
    os.environ["CLEVERAGENTS_AUTO_APPLY_MIGRATIONS"] = "true"

    # Flush the in-memory engine cache so no stale engines leak between
    # scenarios.  In serial mode (coverage_report) every feature shares
    # the same process, so a previous scenario's real-or-fake engine can
    # survive into the next one.
    try:
        from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES

        for _url, engine in list(MEMORY_ENGINES.items()):
            with contextlib.suppress(Exception):
                engine.dispose()
        MEMORY_ENGINES.clear()
    except ImportError:
        pass

    # Clean up any lingering test environment variables from previous tests
    for env_var in [
        "CLEVERAGENTS_MOCK_SHOULD_FAIL",
        "CLEVERAGENTS_MOCK_INVALID_CODE",
        "CLEVERAGENTS_SERVER_MODE",
        *LANGSMITH_ENV_VARS,
    ]:
        if env_var in os.environ:
            del os.environ[env_var]

    # Give each scenario a unique database file so scenarios cannot share
    # persisted state AND parallel subprocesses never collide on the same
    # SQLite file.  Store the paths for cleanup in after_scenario.
    #
    # Features tagged @mock_only use fully mocked services and never touch
    # the database, so skip the temp-file creation for them (~0.5ms each,
    # but the real savings come from not triggering MigrationRunner later).
    context._scenario_db_paths = []
    _is_mock_only = "mock_only" in scenario.effective_tags
    if not _is_mock_only:
        for env_var, prefix in (
            ("CLEVERAGENTS_DATABASE_URL", "cleveragents_"),
            ("CLEVERAGENTS_TEST_DATABASE_URL", "cleveragents_test_"),
        ):
            db_path = tempfile.mktemp(suffix=".db", prefix=prefix)
            os.environ[env_var] = f"sqlite:///{db_path}"
            context._scenario_db_paths.append(db_path)

    # Re-apply mock AI provider after container reset
    try:
        from cleveragents.application.container import override_providers
        from features.mocks.mock_ai_provider import MockAIProvider

        # Override the AI provider with mock for all tests
        mock_provider = MockAIProvider()
        override_providers(ai_provider=mock_provider)
    except ImportError:
        pass  # Container not needed for all tests


def after_scenario(context, scenario):
    """Clean up after each scenario."""
    # Return to original directory first
    if hasattr(context, "original_cwd"):
        os.chdir(context.original_cwd)

    # Run all cleanup handlers in reverse order
    if hasattr(context, "_cleanup_handlers"):
        for handler in reversed(context._cleanup_handlers):
            with contextlib.suppress(Exception):
                handler()
        context._cleanup_handlers = []

    # Clean up any test directories
    if hasattr(context, "test_dir") and context.test_dir:
        try:
            if Path(context.test_dir).exists():
                shutil.rmtree(context.test_dir)
        except Exception:
            pass  # Ignore cleanup errors
        context.test_dir = None

    # Clean up environment variables set during tests
    if hasattr(context, "env_vars_to_clean"):
        for key in context.env_vars_to_clean:
            os.environ.pop(key, None)
        context.env_vars_to_clean = []

    for env_var in [
        *LANGSMITH_ENV_VARS,
        "CLEVERAGENTS_DATABASE_URL",
        "CLEVERAGENTS_TEST_DATABASE_URL",
    ]:
        os.environ.pop(env_var, None)

    # Reset Settings singleton if it was used
    try:
        from cleveragents.config.settings import Settings

        Settings._instance = None
    except ImportError:
        pass

    # Reset global container singleton
    try:
        from cleveragents.application.container import reset_container

        reset_container()
    except ImportError:
        pass

    # Clean up service instances
    for attr in ["context_service", "plan_service", "project_service"]:
        if hasattr(context, attr):
            delattr(context, attr)

    # Clean up any remaining attributes that might hold state
    for attr in ["plan", "plans", "project", "changes", "added_files", "all_plans"]:
        if hasattr(context, attr):
            delattr(context, attr)

    # Restore provider registry env overrides if used in a scenario
    for attr, env_var in [
        ("original_provider_env", "CLEVERAGENTS_DEFAULT_PROVIDER"),
        ("original_model_env", "CLEVERAGENTS_DEFAULT_MODEL"),
    ]:
        if hasattr(context, attr):
            original = getattr(context, attr)
            if original is None:
                os.environ.pop(env_var, None)
            else:
                os.environ[env_var] = original
            delattr(context, attr)

    # Reset provider registry singleton between scenarios
    try:
        from cleveragents.providers.registry import (
            reset_provider_registry as _reset_registry,
        )

        _reset_registry()
    except ImportError:
        pass

    # Clean up in-memory database engine cache to ensure each scenario
    # gets a fresh database. This is critical for tests using :memory: databases
    try:
        from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES

        # Dispose of all cached engines and clear the cache
        for _url, engine in list(MEMORY_ENGINES.items()):
            with contextlib.suppress(Exception):
                engine.dispose()
        MEMORY_ENGINES.clear()
    except ImportError:
        pass

    # Remove per-scenario temp database files (and associated journal/WAL
    # files) now that all engines have been disposed.
    for db_path in getattr(context, "_scenario_db_paths", []):
        for suffix in ("", "-journal", "-wal", "-shm"):
            with contextlib.suppress(OSError):
                os.unlink(db_path + suffix)

    # T6: Remove log handlers attached to the async-cleanup logger by
    # security_async_steps.py so handlers don't accumulate across scenarios.
    import logging

    if hasattr(context, "log_handler"):
        async_logger = logging.getLogger("cleveragents.core.async_cleanup")
        async_logger.removeHandler(context.log_handler)

    # T5: Close event loops left open by security_async_steps.py.
    if hasattr(context, "bridge_loop"):
        with contextlib.suppress(Exception):
            context.bridge_loop.close()