cleveragents-core/tests/features/environment.py

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

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


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_GEMINI_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
    import shutil

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

    # Close async event loop with proper cleanup
    if hasattr(context, "loop"):
        # First, explicitly close any open subprocess transports
        try:
            # Access the loop's selector to find open transports
            if hasattr(context.loop, "_selector") and context.loop._selector:
                # Get all file descriptors that might be subprocess transports
                transports_to_close = []
                if hasattr(context.loop, "_transports"):
                    for fd, transport in context.loop._transports.items():
                        if (
                            hasattr(transport, "__class__")
                            and "Subprocess" in transport.__class__.__name__
                        ):
                            transports_to_close.append(transport)

                # Close subprocess transports
                for transport in transports_to_close:
                    try:
                        if hasattr(transport, "close"):
                            transport.close()
                    except Exception:
                        pass
        except Exception:
            pass

        # Cancel all remaining tasks
        pending = asyncio.all_tasks(context.loop)
        for task in pending:
            task.cancel()

        # Wait for tasks to finish cancelling
        if pending:
            try:
                context.loop.run_until_complete(
                    asyncio.gather(*pending, return_exceptions=True)
                )
            except Exception:
                pass

        # Allow extra time for subprocess transports to cleanup
        try:
            context.loop.run_until_complete(asyncio.sleep(0.2))
        except Exception:
            pass

        # Force garbage collection to help cleanup
        try:
            import gc

            gc.collect()
        except Exception:
            pass

        # Close the loop
        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

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

    # Set up test context for InlineYAMLJinja tests - not needed for simplified tests
    # if "InlineYAMLJinja" in scenario.feature.name or "inline_yaml_jinja" in str(scenario.feature.filename):
    #     from tests.features.steps.inline_yaml_jinja_coverage_steps import TestContext
    #     context.test_context = TestContext()


def after_scenario(context, scenario):
    """Clean up after each scenario."""
    # Dispose of app if created
    if hasattr(context, "app") and context.app:
        context.app.dispose()

    # Clear any running tasks
    if hasattr(context, "loop"):
        # First close any subprocess transports
        try:
            if hasattr(context.loop, "_transports"):
                transports_to_close = []
                for fd, transport in list(context.loop._transports.items()):
                    if (
                        hasattr(transport, "__class__")
                        and "Subprocess" in transport.__class__.__name__
                    ):
                        transports_to_close.append(transport)

                for transport in transports_to_close:
                    try:
                        if hasattr(transport, "close") and not transport.is_closing():
                            transport.close()
                    except Exception:
                        pass
        except Exception:
            pass

        # Cancel all remaining tasks
        pending = asyncio.all_tasks(context.loop)
        for task in pending:
            task.cancel()

        if pending:
            try:
                context.loop.run_until_complete(
                    asyncio.gather(*pending, return_exceptions=True)
                )
            except Exception:
                pass

        # Allow more time for subprocess transports to cleanup
        try:
            context.loop.run_until_complete(asyncio.sleep(0.05))
        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()
                    if tasks:
                        # Run the loop briefly to let tasks complete cancellation
                        try:
                            loop.run_until_complete(
                                asyncio.gather(*tasks, return_exceptions=True)
                            )
                        except:
                            pass
                    # Allow time for subprocess transports to cleanup
                    try:
                        loop.run_until_complete(asyncio.sleep(0.05))
                    except:
                        pass
                    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:
                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()
                        if tasks:
                            try:
                                loop.run_until_complete(
                                    asyncio.gather(*tasks, return_exceptions=True)
                                )
                            except:
                                pass
                        # Allow time for subprocess transports to cleanup
                        try:
                            loop.run_until_complete(asyncio.sleep(0.05))
                        except:
                            pass
                        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 test context for InlineYAMLJinja tests - not needed for simplified tests
    # if hasattr(context, 'test_context'):
    #     context.test_context.cleanup()