cleveractors-core/features/environment.py

"""
BDD test environment setup for reactive CleverAgents.
"""

import asyncio
import os
import shutil
import sys
import tempfile
from pathlib import Path

# Increase recursion limit to handle deep import chains in langchain/torch/transformers
sys.setrecursionlimit(5000)


def before_all(context):
    """Set up test environment before all tests."""
    context.temp_dir = Path(tempfile.mkdtemp())

    # Set dummy API keys for testing
    os.environ["OPENAI_API_KEY"] = "test-key-openai"
    os.environ["ANTHROPIC_API_KEY"] = "test-key-anthropic"
    os.environ["GOOGLE_API_KEY"] = "test-key-google"

    # Create async event loop for tests
    context.loop = asyncio.new_event_loop()
    asyncio.set_event_loop(context.loop)


def after_all(context):
    """Clean up test environment after all tests."""
    # Clean up temp directory

    shutil.rmtree(context.temp_dir, ignore_errors=True)

    # Close async event loop with proper cleanup
    if hasattr(context, "loop"):
        # Cancel all remaining tasks
        try:
            pending = asyncio.all_tasks(context.loop)
            for task in pending:
                task.cancel()
        except Exception:
            pass

        # Close the loop without waiting for tasks
        try:
            context.loop.close()
        except Exception:
            pass


def before_scenario(context, scenario):
    """Set up before each scenario."""
    context.config_files = []
    context.app = None
    context.result = None
    context.error = None
    context.unsafe = False
    context._active_patches = []  # Defensive init — after_scenario cleanup iterates this
    context._llm_should_fail = (
        False  # Reset from prior scenarios that configure LLM failure mode
    )

    # Create scenario-specific temp directory
    context.scenario_temp = (
        context.temp_dir
        / f"scenario_{scenario.name.replace(' ', '_').replace('/', '_')}"
    )
    context.scenario_temp.mkdir(exist_ok=True)

    # Create a fresh event loop per scenario to avoid "event loop is closed" errors
    # from async_step decorators when the loop has been corrupted by prior cleanup
    try:
        old_loop = asyncio.get_event_loop()
        if old_loop.is_closed():
            new_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(new_loop)
            context.loop = new_loop
    except RuntimeError:
        new_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(new_loop)
        context.loop = new_loop

    # Always ensure API keys are set (some steps clear them and never restore)
    os.environ.setdefault("OPENAI_API_KEY", "test-key-openai")
    os.environ.setdefault("ANTHROPIC_API_KEY", "test-key-anthropic")
    os.environ.setdefault("GOOGLE_API_KEY", "test-key-google")


def after_scenario(context, scenario):
    """Clean up after each scenario."""
    # Clean up bridge resources if present
    if hasattr(context, "bridge") and context.bridge:
        try:
            context.bridge.cleanup()
        except Exception:
            pass

    # Dispose of app if created
    if hasattr(context, "app") and context.app:
        try:
            # dispose() is async, so we need to run it in the event loop
            if hasattr(context, "loop") and context.loop:
                context.loop.run_until_complete(context.app.dispose())
            else:
                asyncio.run(context.app.dispose())
        except Exception:
            pass

    # Cancel any remaining tasks without waiting
    if hasattr(context, "loop"):
        try:
            pending = asyncio.all_tasks(context.loop)
            for task in pending:
                task.cancel()
            # Drain cancellations so tasks get cleaned up and do not
            # accumulate on the shared loop across scenarios.
            _drain = context.loop.create_task(asyncio.sleep(0))
            context.loop.run_until_complete(_drain)
        except Exception:
            pass

    # Clean up event loops created by LangGraph tests
    if hasattr(context, "event_loops"):
        for loop in context.event_loops:
            try:
                if hasattr(loop, "is_closed") and not loop.is_closed():
                    # Cancel all tasks in this loop
                    tasks = asyncio.all_tasks(loop)
                    for task in tasks:
                        task.cancel()
                    loop.close()
            except Exception:
                pass

    # Clean up temp directories from LangGraph tests
    if hasattr(context, "temp_dirs"):
        for temp_dir in context.temp_dirs:
            try:
                shutil.rmtree(temp_dir)
            except Exception:
                pass

    # Clean up graphs and their schedulers
    if hasattr(context, "graphs"):
        for graph in context.graphs:
            try:
                if hasattr(graph, "scheduler") and hasattr(graph.scheduler, "_loop"):
                    loop = graph.scheduler._loop
                    if hasattr(loop, "is_closed") and not loop.is_closed():
                        # Cancel all tasks
                        tasks = asyncio.all_tasks(loop)
                        for task in tasks:
                            task.cancel()
                        loop.close()
            except Exception:
                pass

    # Clean up template loader test files
    if "template_loaders_coverage" in str(scenario.feature.filename):
        # fmt: off
        from tests.features.steps.template_loaders_coverage_steps import cleanup_test_files  # isort: skip
        # fmt: on

        cleanup_test_files(context)

    # Clean up files tracked for cleanup in tool_agent tests
    if hasattr(context, "__dict__") and "_cleanup_files" in context.__dict__:
        for filepath in context.__dict__["_cleanup_files"]:
            try:
                if os.path.exists(filepath):
                    os.remove(filepath)
            except Exception:
                pass
        context.__dict__["_cleanup_files"] = []

    # Stop any env_patch that wasn't stopped by a Then step (m-5).
    if hasattr(context, "env_patch"):
        try:
            context.env_patch.stop()
        except Exception:
            pass

    # Stop any _prune_patch from pruning model tests.
    if hasattr(context, "_prune_patch") and context._prune_patch is not None:
        try:
            context._prune_patch.stop()
        except Exception:
            pass
        context._prune_patch = None

    # Remove log-capture handler installed by cleanup warning-log tests.
    if hasattr(context, "_cleanup_log_handler"):
        from cleveractors.agents.llm import logger as llm_logger

        try:
            llm_logger.removeHandler(context._cleanup_log_handler)
        except Exception:
            pass
        if hasattr(context, "_cleanup_log_original_level"):
            try:
                llm_logger.setLevel(context._cleanup_log_original_level)
            except Exception:
                pass

    # Stop any active patches used by test scenarios.
    if hasattr(context, "_active_patches"):
        for p in reversed(context._active_patches):
            try:
                p.stop()
            except Exception:
                pass
        context._active_patches = []

    # Clean up scenario temp directory
    if hasattr(context, "scenario_temp") and context.scenario_temp.exists():
        try:
            shutil.rmtree(context.scenario_temp, ignore_errors=True)
        except Exception:
            pass

    # Clean up _tmpdir from LocalPackageStore tests
    if hasattr(context, "_tmpdir"):
        try:
            context._tmpdir.cleanup()
        except Exception:
            pass

    # Reset budget ContextVars set by test steps (issue #76 test hygiene).
    # Without this, budget ContextVars set in llm_agent_tool_loop scenarios
    # leak across scenarios and can cause order-dependent flakiness
    # (Graa review PR #80).
    if hasattr(context, "_budget_cost_token"):
        from cleveractors.agents.retry import (
            current_accumulated_cost,
            current_max_cost_usd,
            current_pricing,
        )

        try:
            current_accumulated_cost.reset(context._budget_cost_token)
            if getattr(context, "_budget_max_token", None) is not None:
                current_max_cost_usd.reset(context._budget_max_token)
            if getattr(context, "_budget_pricing_token", None) is not None:
                current_pricing.reset(context._budget_pricing_token)
        except Exception:
            pass

    # Restore monkey-patched validation constants
    if hasattr(context, "_original_max_dfs"):
        import cleveractors.validation._limits as _limits

        _limits._MAX_DFS_STEPS = context._original_max_dfs
    if hasattr(context, "_original_max_subgraph"):
        import cleveractors.validation._limits as _limits

        _limits._MAX_SUBGRAPH_STEPS = context._original_max_subgraph

    # Clean up log-capture handler from domain-pattern warning scenarios.
    if hasattr(context, "_domain_warning_handler"):
        try:
            from cleveractors.agents.llm_providers import logger as llm_providers_logger

            llm_providers_logger.removeHandler(context._domain_warning_handler)
        except Exception:
            pass
    if hasattr(context, "_domain_warning_original_level"):
        try:
            from cleveractors.agents.llm_providers import logger as llm_providers_logger

            llm_providers_logger.setLevel(context._domain_warning_original_level)
        except Exception:
            pass

    # Clean up log-capture handler from llm_client warning scenarios.
    if hasattr(context, "_llm_client_warning_handler"):
        try:
            from cleveractors.agents.llm_client import logger as llm_client_logger

            llm_client_logger.removeHandler(context._llm_client_warning_handler)
        except Exception:
            pass
    if hasattr(context, "_llm_client_warning_original_level"):
        try:
            from cleveractors.agents.llm_client import logger as llm_client_logger

            llm_client_logger.setLevel(context._llm_client_warning_original_level)
        except Exception:
            pass

    # Clean up log-capture handler from _execute_llm cleanup scenarios.
    if hasattr(context, "_llm_cleanup_handler"):
        try:
            from cleveractors.runtime_dispatch import logger as dispatch_logger

            dispatch_logger.removeHandler(context._llm_cleanup_handler)
        except Exception:
            pass
    if hasattr(context, "_llm_cleanup_original_level"):
        try:
            from cleveractors.runtime_dispatch import logger as dispatch_logger

            dispatch_logger.setLevel(context._llm_cleanup_original_level)
        except Exception:
            pass

    # Clean up log-capture handler from _execute_graph cleanup scenarios.
    if hasattr(context, "_graph_cleanup_handler"):
        try:
            from cleveractors.runtime_dispatch import logger as dispatch_logger

            dispatch_logger.removeHandler(context._graph_cleanup_handler)
        except Exception:
            pass
    if hasattr(context, "_graph_cleanup_original_level"):
        try:
            from cleveractors.runtime_dispatch import logger as dispatch_logger

            dispatch_logger.setLevel(context._graph_cleanup_original_level)
        except Exception:
            pass