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

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


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",
        *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)