conftest.py

"""Allows you to use `pytest docs` to run the examples.

To run notebooks, use: uv run --with 'mcp' pytest --nbmake docs/examples/notebooks/
"""

import ast
import os
import pathlib
import subprocess
import sys

import pytest

# Cached result of system capability detection (None = not yet computed)
_capabilities_cache: dict | None = None


def get_system_capabilities():
    """Lazy load system capabilities from test/conftest.py, cached after first call."""
    global _capabilities_cache

    if _capabilities_cache is not None:
        return _capabilities_cache

    # Add test directory to path to enable import
    _test_dir = pathlib.Path(__file__).parent.parent.parent / "test"
    _test_dir_abs = _test_dir.resolve()
    if str(_test_dir_abs) not in sys.path:
        sys.path.insert(0, str(_test_dir_abs))

    try:
        # Import with explicit module name to avoid conflicts
        import importlib.util

        spec = importlib.util.spec_from_file_location(
            "test_conftest", _test_dir_abs / "conftest.py"
        )
        if spec and spec.loader:
            test_conftest = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(test_conftest)
            _capabilities_cache = test_conftest.get_system_capabilities()
            return _capabilities_cache
        else:
            raise ImportError("Could not load test/conftest.py")
    except (ImportError, AttributeError) as e:
        # Fallback if test/conftest.py not available
        import warnings

        warnings.warn(
            f"Could not import get_system_capabilities from test/conftest.py: {e}. Heavy RAM tests will NOT be skipped!"
        )

        _capabilities_cache = {
            "has_gpu": False,
            "gpu_memory_gb": 0,
            "ram_gb": 0,
            "has_api_keys": {},
            "has_ollama": False,
        }
        return _capabilities_cache


examples_to_skip: dict[str, str] = {}


def _extract_markers_from_file(file_path):
    """Extract pytest markers from comment in file without parsing Python.

    Looks for lines like: # pytest: marker1, marker2, marker3
    Returns list of marker names.
    """
    try:
        with open(file_path) as f:
            for line in f:
                line = line.strip()
                if line.startswith("# pytest:"):
                    marker_text = line[9:].strip()  # Remove "# pytest:"
                    return [m.strip() for m in marker_text.split(",") if m.strip()]
                # Stop after first few lines (markers should be at top)
                if (
                    len(line) > 0
                    and not line.startswith("#")
                    and not line.startswith('"""')
                ):
                    break
    except Exception:
        pass
    return []


def _should_skip_collection(markers):
    """Check if example should be skipped during collection based on markers.

    Returns (should_skip, reason) tuple.
    """
    if not markers:
        return False, None

    # Skip tests marked with skip_always
    if "skip_always" in markers:
        return True, "Example marked to always skip (skip_always marker)"

    try:
        capabilities = get_system_capabilities()
    except Exception:
        # If we can't get capabilities, don't skip (fail open)
        return False, None

    gh_run = int(os.environ.get("CICD", 0))

    # Skip qualitative tests in CI
    if "qualitative" in markers and gh_run == 1:
        return True, "Skipping qualitative test in CI (CICD=1)"

    # Explicitly skip if 'skip' marker is present
    if "skip" in markers:
        return True, "Example marked with skip marker"

    # Skip slow tests if SKIP_SLOW=1 environment variable is set
    if "slow" in markers and int(os.environ.get("SKIP_SLOW", 0)) == 1:
        return True, "Skipping slow test (SKIP_SLOW=1)"

    # Skip tests requiring GPU if not available
    if "huggingface" in markers or "vllm" in markers:
        if not capabilities["has_gpu"]:
            return True, "GPU not available"

    # Skip tests requiring Ollama if not available
    if "ollama" in markers:
        if not capabilities["has_ollama"]:
            return True, "Ollama not available (port 11434 not listening)"

    # Skip tests requiring API keys
    if "watsonx" in markers:
        if not capabilities["has_api_keys"].get("watsonx"):
            return True, "Watsonx API credentials not found"
    if "openai" in markers:
        if not capabilities["has_api_keys"].get("openai"):
            return True, "OpenAI API key not found"

    return False, None


def pytest_addoption(parser):
    """Add command-line options for skipping capability checks.

    These match the options in test/conftest.py to provide consistent behavior.
    Only adds options if they don't already exist (to avoid conflicts when both
    test/ and docs/ conftest files are loaded).
    """

    # Helper to safely add option only if it doesn't exist
    def add_option_safe(option_name, **kwargs):
        try:
            parser.addoption(option_name, **kwargs)
        except ValueError:
            # Option already exists (likely from test/conftest.py)
            pass

    add_option_safe(
        "--ignore-gpu-check",
        action="store_true",
        default=False,
        help="Ignore GPU requirement checks (examples may fail without GPU)",
    )
    add_option_safe(
        "--ignore-ram-check",
        action="store_true",
        default=False,
        help="Ignore RAM requirement checks (examples may fail with insufficient RAM)",
    )
    add_option_safe(
        "--ignore-ollama-check",
        action="store_true",
        default=False,
        help="Ignore Ollama availability checks (examples will fail if Ollama not running)",
    )
    add_option_safe(
        "--ignore-api-key-check",
        action="store_true",
        default=False,
        help="Ignore API key checks (examples will fail without valid API keys)",
    )
    add_option_safe(
        "--ignore-all-checks",
        action="store_true",
        default=False,
        help="Ignore all requirement checks (GPU, RAM, Ollama, API keys)",
    )


def _collect_vllm_example_files(session) -> list[str]:
    """Collect all example files that have vLLM marker.

    Returns list of file paths.
    """
    vllm_files = set()

    for item in session.items:
        # Check if this is an ExampleItem with vllm marker
        if hasattr(item, "path"):
            file_path = str(item.path)
            # Check if file has vllm marker
            if file_path.endswith(".py"):
                markers = _extract_markers_from_file(file_path)
                if "vllm" in markers:
                    vllm_files.add(file_path)

    return sorted(vllm_files)


def _run_vllm_examples_isolated(session, vllm_files: list[str]) -> int:
    """Run vLLM example files in separate processes for GPU memory isolation.

    Returns exit code (0 = all passed, 1 = any failed).
    """
    print("\n" + "=" * 70)
    print("vLLM Process Isolation Active (Examples)")
    print("=" * 70)
    print(f"Running {len(vllm_files)} vLLM example(s) in separate processes")
    print("to ensure GPU memory is fully released between examples.\n")

    # Set environment variables for vLLM
    env = os.environ.copy()
    env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

    all_passed = True

    for i, file_path in enumerate(vllm_files, 1):
        print(f"\n[{i}/{len(vllm_files)}] Running: {file_path}")
        print("-" * 70)

        # Run example directly with Python
        cmd = [sys.executable, file_path]

        result = subprocess.run(cmd, env=env)

        if result.returncode != 0:
            all_passed = False
            print(f"✗ Example failed: {file_path}")
        else:
            print(f"✓ Example passed: {file_path}")

    print("\n" + "=" * 70)
    if all_passed:
        print("All vLLM examples passed!")
    else:
        print("Some vLLM examples failed.")
    print("=" * 70 + "\n")

    return 0 if all_passed else 1


def pytest_collection_finish(session):
    """After collection, check if we need vLLM process isolation for examples.

    If vLLM examples are collected and there are multiple files,
    run them in separate processes and exit.
    """
    # Only check for examples in docs/examples
    if not any(
        "docs" in str(item.path) and "examples" in str(item.path)
        for item in session.items
    ):
        return

    # Collect vLLM example files
    vllm_files = _collect_vllm_example_files(session)

    # Only use process isolation if multiple vLLM examples
    if len(vllm_files) <= 1:
        return

    # Run examples in isolation
    exit_code = _run_vllm_examples_isolated(session, vllm_files)

    # Clear collected items so pytest doesn't run them again
    session.items.clear()

    # Exit with appropriate code
    pytest.exit("vLLM examples completed in isolated processes", returncode=exit_code)


def pytest_terminal_summary(terminalreporter, exitstatus, config):
    if examples_to_skip:
        terminalreporter.ensure_newline()
        terminalreporter.section("Skipped Examples", sep="=", blue=True, bold=True)
        terminalreporter.line(
            "The following examples were skipped during collection:\n"
        )
        for filepath, reason in examples_to_skip.items():
            terminalreporter.line(f"  • {pathlib.Path(filepath).name}: {reason}")


def pytest_pycollect_makemodule(module_path, parent):
    """Intercepts Module creation to skip files before import.

    Only fires for files matching python_files (default test_*.py) during
    directory traversal, or for any file specified directly on the command
    line. Returning a SkippedFile prevents pytest from importing the file,
    which is necessary when files contain unavailable dependencies.

    Args:
        module_path: pathlib.Path to the module
        parent: Parent collector node
    """
    file_path = module_path

    # Limit scope to docs/examples directory
    if "docs" not in file_path.parts or "examples" not in file_path.parts:
        return None

    if file_path.name == "conftest.py":
        return None

    # Initialize capabilities cache if needed
    config = parent.config
    if not hasattr(config, "_example_capabilities"):
        config._example_capabilities = get_system_capabilities()

    # Check manual skip list
    if str(file_path) in examples_to_skip:
        return SkippedFile.from_parent(parent, path=file_path)

    # Extract and evaluate markers
    markers = _extract_markers_from_file(file_path)

    if not markers:
        return SkippedFile.from_parent(parent, path=file_path)

    should_skip, _reason = _should_skip_collection(markers)

    if should_skip:
        # Prevent import by returning custom collector
        return SkippedFile.from_parent(parent, path=file_path)

    # Return ExampleModule so pytest never falls through to its default Module
    # collector (which would import the file directly). Import errors are
    # instead caught at runtime in ExampleItem.runtest() and converted to skips.
    return ExampleModule.from_parent(parent, path=file_path)


def pytest_ignore_collect(collection_path, config):
    """Ignore files before pytest even tries to parse them.

    This is called BEFORE pytest_collect_file, so we can prevent
    heavy files from being parsed at all.

    NOTE: This hook is only called during directory traversal, not for
    directly specified files. The pytest_pycollect_makemodule hook handles
    both cases.
    """
    # Skip conftest.py itself - it's not a test
    if collection_path.name == "conftest.py":
        return True

    # Convert to absolute path to check if it's in docs/examples
    # (pytest may pass relative paths)
    abs_path = collection_path.resolve()

    # Only check Python files in docs/examples
    if (
        collection_path.suffix == ".py"
        and "docs" in abs_path.parts
        and "examples" in abs_path.parts
    ):
        # Skip files in the manual skip list
        if str(collection_path) in examples_to_skip:
            return True

        # Extract markers and check if we should skip
        try:
            markers = _extract_markers_from_file(collection_path)
            # No markers → not a runnable example (e.g. __init__.py, helpers)
            if not markers:
                return True
            should_skip, reason = _should_skip_collection(markers)
            if should_skip and reason:
                # Add to skip list with reason for terminal summary
                examples_to_skip[str(collection_path)] = reason
                # Return True to ignore this file completely
                return True
        except Exception as e:
            # Log the error but don't skip - let pytest handle it
            import sys

            print(
                f"WARNING: Error checking markers for {collection_path}: {e}",
                file=sys.stderr,
            )

    return False


def pytest_collect_file(parent: pytest.Dir, file_path: pathlib.PosixPath):
    """Provide an explicit collector for example files in docs/examples/."""
    if (
        file_path.suffix == ".py"
        and "docs" in file_path.parts
        and "examples" in file_path.parts
    ):
        # Directly-specified files are handled by pytest_pycollect_makemodule —
        # only provide an explicit collector during directory traversal.
        if parent.session.isinitpath(file_path):
            return None

        # Already flagged for skipping (missing system capability)
        if str(file_path) in examples_to_skip:
            return

        # Check markers — no markers means not a runnable example.
        # _extract_markers_from_file is self-contained (returns [] on error),
        # so no try/except needed here.
        markers = _extract_markers_from_file(file_path)
        if not markers:
            return None
        should_skip, _reason = _should_skip_collection(markers)
        if should_skip:
            return SkippedFile.from_parent(parent, path=file_path)

        # pytest_pycollect_makemodule only fires for files matching python_files
        # (test_*.py) — examples need an explicit collector for directory traversal.
        return ExampleModule.from_parent(parent, path=file_path)


class SkippedFile(pytest.File):
    """A dummy collector for skipped files to prevent default import.

    This collector is returned by pytest_pycollect_makemodule and pytest_collect_file
    when a file should be skipped based on markers or system capabilities.

    By returning this custom collector instead of None, we prevent pytest from
    falling back to its default Module collector which would import the file.
    The collect() method returns an empty list, so no tests are collected.
    """

    def __init__(self, **kwargs):
        # Extract reason if provided, otherwise use default
        self.skip_reason = kwargs.pop("reason", "File skipped based on markers")
        super().__init__(**kwargs)

    def collect(self):
        # Return empty list - no tests to collect from this file
        return []


class ExampleFile(pytest.File):
    def collect(self):
        return [ExampleItem.from_parent(self, name=self.name)]


class ExampleModule(pytest.Module):
    """Module stand-in that routes to ExampleItem without importing the file.

    Returned by pytest_pycollect_makemodule to prevent pytest's default
    Module collector from importing the file directly (which would crash on
    missing optional deps before ExampleItem.runtest() can catch them).
    """

    def collect(self):
        return [ExampleItem.from_parent(self, name=self.path.name)]


class ExampleItem(pytest.Item):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def runtest(self):
        import os
        import pathlib

        repo_root = str(pathlib.Path(__file__).parent.parent.parent.resolve())
        env = os.environ.copy()
        existing_pythonpath = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = (
            f"{existing_pythonpath}{os.pathsep}{repo_root}"
            if existing_pythonpath
            else repo_root
        )

        process = subprocess.Popen(
            [sys.executable, self.path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,  # Enable line-buffering
            env=env,
        )

        # Capture stdout output and output it so it behaves like a regular test with -s.
        stdout_lines = []
        if process.stdout is not None:
            for line in process.stdout:
                sys.stdout.write(line)
                sys.stdout.flush()  # Ensure the output is printed immediately
                stdout_lines.append(line)
            process.stdout.close()

        retcode = process.wait()

        # Capture stderr output.
        stderr = ""
        if process.stderr is not None:
            stderr = process.stderr.read()

        if retcode != 0:
            # Check if this is a pytest.skip() call (indicated by "Skipped:" in stderr)
            if "Skipped:" in stderr or "_pytest.outcomes.Skipped" in stderr:
                # Extract skip reason from stderr
                skip_reason = "Example skipped"
                for line in stderr.split("\n"):
                    if line.startswith("Skipped:"):
                        skip_reason = line.replace("Skipped:", "").strip()
                        break
                pytest.skip(skip_reason)
            elif "ModuleNotFoundError" in stderr or "ImportError" in stderr:
                # Missing optional dependency — skip rather than fail so the
                # suite stays green without every optional package installed.
                reason = "optional dependency not installed"
                for line in stderr.split("\n"):
                    if "ModuleNotFoundError" in line or "ImportError" in line:
                        reason = line.strip()
                        break
                pytest.skip(reason)
            else:
                raise ExampleTestException(
                    f"Example failed with exit code {retcode}.\nStderr: {stderr}\n"
                )

    def repr_failure(self, excinfo, style=None):
        """Called when self.runtest() raises an exception."""
        if isinstance(excinfo.value, ExampleTestException):
            return str(excinfo.value)

        return super().repr_failure(excinfo)

    def reportinfo(self):
        return self.path, 0, f"usecase: {self.name}"


class ExampleTestException(Exception):
    """Custom exception for error reporting."""


def pytest_runtest_setup(item):
    """Apply skip logic to ExampleItem objects based on system capabilities.

    This ensures examples respect the same capability checks as regular tests
    (RAM, GPU, Ollama, API keys, etc.).
    """
    if not isinstance(item, ExampleItem):
        return

    # Check for explicit skip marker first
    if item.get_closest_marker("skip"):
        pytest.skip("Example marked with skip marker")

    # Get system capabilities
    capabilities = get_system_capabilities()

    # Get gh_run status (CI environment)
    gh_run = int(os.environ.get("CICD", 0))

    # Get config options from CLI (matching test/conftest.py behavior)
    config = item.config
    ignore_all = config.getoption("--ignore-all-checks", default=False)
    ignore_gpu = config.getoption("--ignore-gpu-check", default=False) or ignore_all
    ignore_ollama = (
        config.getoption("--ignore-ollama-check", default=False) or ignore_all
    )
    ignore_api_key = (
        config.getoption("--ignore-api-key-check", default=False) or ignore_all
    )

    # Skip qualitative tests in CI
    if item.get_closest_marker("qualitative") and gh_run == 1:
        pytest.skip(
            reason="Skipping qualitative test: got env variable CICD == 1. Used only in gh workflows."
        )

    # Skip tests requiring GPU if not available
    if (
        item.get_closest_marker("huggingface") or item.get_closest_marker("vllm")
    ) and not ignore_gpu:
        if not capabilities["has_gpu"]:
            pytest.skip("Skipping test: GPU not available")

    # Backend-specific skipping
    if item.get_closest_marker("watsonx") and not ignore_api_key:
        if not capabilities["has_api_keys"].get("watsonx"):
            pytest.skip(
                "Skipping test: Watsonx API credentials not found in environment"
            )

    if item.get_closest_marker("vllm") and not ignore_gpu:
        if not capabilities["has_gpu"]:
            pytest.skip("Skipping test: vLLM requires GPU")

    if item.get_closest_marker("ollama") and not ignore_ollama:
        if not capabilities["has_ollama"]:
            pytest.skip(
                "Skipping test: Ollama not available (port 11434 not listening)"
            )


def pytest_runtest_teardown(item, nextitem):
    """Evict Ollama models after each ollama-marked example.

    Examples run as subprocesses, so Ollama's default keep_alive keeps
    models resident after exit. Evict after every example to prevent
    heavyweight models from starving subsequent examples of memory (#798).
    """
    if not isinstance(item, ExampleItem):
        return
    if not item.get_closest_marker("ollama"):
        return

    _evict_ollama_models()


def _evict_ollama_models() -> None:
    """Evict all currently loaded Ollama models (best-effort)."""
    import requests

    host = os.environ.get("OLLAMA_HOST", "127.0.0.1")
    if ":" in host:
        host, port = host.rsplit(":", 1)
    else:
        port = os.environ.get("OLLAMA_PORT", "11434")

    if host == "0.0.0.0":
        host = "127.0.0.1"

    base_url = f"http://{host}:{port}"

    try:
        resp = requests.get(f"{base_url}/api/ps", timeout=5)
        resp.raise_for_status()
        loaded = resp.json().get("models", [])
    except Exception:
        return

    for entry in loaded:
        model_name = entry.get("name") or entry.get("model", "unknown")
        try:
            requests.post(
                f"{base_url}/api/generate",
                json={"model": model_name, "keep_alive": 0},
                timeout=10,
            )
            print(f"ollama-evict: evicted {model_name}", file=sys.stderr)
        except Exception:
            pass


def pytest_collection_modifyitems(items):
    """Apply markers from example files to ExampleItem objects.

    Parses comment-based markers from example files in the format:
        # pytest: marker1, marker2, marker3

    This keeps examples clean while allowing intelligent test skipping.
    """
    for item in items:
        if isinstance(item, ExampleItem):
            # Read the file and look for comment-based markers
            try:
                with open(item.path) as f:
                    for line in f:
                        line = line.strip()
                        # Look for comment-based marker line
                        if line.startswith("# pytest:"):
                            # Extract markers after "# pytest:"
                            marker_text = line[9:].strip()  # Remove "# pytest:"
                            markers = [m.strip() for m in marker_text.split(",")]
                            for marker_name in markers:
                                if marker_name:  # Skip empty strings
                                    item.add_marker(getattr(pytest.mark, marker_name))
                            break  # Only process first pytest comment line
            except Exception:
                # If we can't parse the file, skip marker application
                pass