cleveragents-core/features/environment.py

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

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

from behave.model import Scenario
from behave.model_core import Status

_tdd_logger = logging.getLogger("behave.tdd_expected_fail")

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)

    # Clear devcontainer lifecycle registry between scenarios to prevent
    # test pollution from in-memory lifecycle trackers and health check
    # threads left by previous scenarios.
    try:
        from cleveragents.resource.handlers.devcontainer import (
            clear_lifecycle_registry,
        )

        clear_lifecycle_registry()
    except ImportError:
        pass  # Handler not available in all environments

    # 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


# ---------------------------------------------------------------------------
# TDD expected-fail tag handler (see CONTRIBUTING.md §TDD Bug Test Tags)
# ---------------------------------------------------------------------------

_TDD_BUG_N_RE = re.compile(r"^tdd_bug_\d+$")


def handle_tdd_expected_fail(scenario: Scenario) -> None:
    """Validate TDD tags and invert pass/fail for ``@tdd_expected_fail``.

    Tag validation (per CONTRIBUTING.md — unconditional):

    * Any scenario with ``@tdd_bug_<N>`` **must** also carry ``@tdd_bug``.
      Missing ``@tdd_bug`` causes the scenario to fail unconditionally.

    Status inversion (only when ``@tdd_expected_fail`` is present):

    * ``@tdd_expected_fail`` additionally requires ``@tdd_bug`` **and** at
      least one ``@tdd_bug_<N>`` tag.
    * **failed → passed** — the bug still triggers, which is expected.
    * **passed → failed** — the bug was fixed but the tag was not removed;
      this is an error that must be caught.
    """
    tags = set(scenario.effective_tags)

    # --- unconditional tag validation (F5) --------------------------------
    has_tdd_bug_n = any(_TDD_BUG_N_RE.match(t) for t in tags)
    if has_tdd_bug_n and "tdd_bug" not in tags:
        scenario.set_status(Status.failed)
        sys.stderr.write(
            f"TDD TAG ERROR: {scenario.name!r} — @tdd_bug_<N> requires @tdd_bug tag\n"
        )
        return

    if "tdd_expected_fail" not in tags:
        return

    # --- @tdd_expected_fail tag validation --------------------------------
    if "tdd_bug" not in tags:
        scenario.set_status(Status.failed)
        sys.stderr.write(
            f"TDD TAG ERROR: {scenario.name!r} — "
            "@tdd_expected_fail requires @tdd_bug tag\n"
        )
        return

    if not has_tdd_bug_n:
        scenario.set_status(Status.failed)
        sys.stderr.write(
            f"TDD TAG ERROR: {scenario.name!r} — "
            "@tdd_expected_fail requires at least one @tdd_bug_<N> tag\n"
        )
        return

    # --- status inversion -------------------------------------------------
    if scenario.status == Status.failed:
        # Log original failure details before inverting so CI logs show what
        # actually failed (N3 review finding).
        for step in scenario.steps:
            if step.status == Status.failed:
                _tdd_logger.info(
                    "TDD inversion: %s — step %r failed: %s",
                    scenario.name,
                    step.name,
                    step.error_message or "(no message)",
                )

        # Bug still present — expected.  Mark scenario and its failed/skipped
        # steps as passed so that summary counts are accurate.
        scenario.clear_status()
        scenario.set_status(Status.passed)
        for step in scenario.steps:
            if step.status in (Status.failed, Status.skipped):
                step.status = Status.passed
    elif scenario.status == Status.passed:
        # Bug was fixed but @tdd_expected_fail was not removed — error.
        scenario.set_status(Status.failed)
        sys.stderr.write(
            f"TDD TAG ERROR: {scenario.name!r} — "
            "scenario passed but still carries @tdd_expected_fail; "
            "remove the tag now that the bug is fixed\n"
        )


def after_scenario(context, scenario):
    """Clean up after each scenario."""
    # Handle TDD expected-fail inversion BEFORE cleanup (status is already set
    # by step execution; cleanup does not change it).
    handle_tdd_expected_fail(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()