cleveragents-core/features/environment.py

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

import contextlib
import os
import shutil
import sys
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")
    os.environ.setdefault("CLEVERAGENTS_DATABASE_URL", "sqlite:///cleveragents.db")
    os.environ.setdefault(
        "CLEVERAGENTS_TEST_DATABASE_URL", "sqlite:///cleveragents_test.db"
    )
    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]

    # Ensure each scenario starts with a fresh database URL so tests
    # don't share persisted project state across runs
    for env_var in ["CLEVERAGENTS_DATABASE_URL", "CLEVERAGENTS_TEST_DATABASE_URL"]:
        os.environ.pop(env_var, None)

    # 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