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

    # --- 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).
    _install_template_db_patch()


def _install_template_db_patch() -> None:
    """Monkey-patch MigrationRunner to copy a template DB for fresh SQLite files."""
    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 when creating per-scenario temp DBs.
    _SCENARIO_DB_PREFIXES = ("cleveragents_", "cleveragents_test_")

    def _fast_init_or_upgrade(self: Any, **kwargs: Any) -> None:
        """Copy the template DB instead of running Alembic migrations.

        Falls through to the original method for:
        - Non-SQLite databases
        - In-memory SQLite databases (:memory:)
        - SQLite files that already exist on disk
        - SQLite files whose basename doesn't match the before_scenario
          temp-file naming pattern (cleveragents_* / cleveragents_test_*)
        """
        db_url: str = getattr(self, "database_url", "")

        # Only intercept file-based SQLite URLs
        if not db_url.startswith("sqlite") or ":memory:" in db_url:
            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 (fresh scenario)
        if db_path.exists():
            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 = {}

    # 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.
    context._scenario_db_paths = []
    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()