litellm/tests/integration/conftest.py

from __future__ import annotations

import hashlib
import json
import os
import re
from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence
from contextlib import ExitStack, contextmanager
from copy import copy
from functools import partial
from importlib.metadata import version
from pathlib import Path
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable

import httpx
import pytest
import pytest_asyncio
from execnet import XSpec
from pydantic import TypeAdapter
from redis import Redis

from tests.integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES
from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups
from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState
from tests.integration._support.provider import SharedProvider, shared_provider
from tests.integration._support.routing import RoutingPlugin
from tests.integration._support.sdk_models import build_sdk_models
from tests.integration._support.workers import gateway_profile, start_worker_stacks, worker_stack
from tests.integration.run import GITHUB_FILES

if TYPE_CHECKING:
    from tests.integration._support.actors import ActorSet, GatewayClient, ManagementDatabase, ScratchResources

COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
INVENTORY: Final = pytest.StashKey[tuple[str, ...]]()
REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]()
WORKER_INVENTORY: Final = TypeAdapter(tuple[str, ...])
SHARD_PATH: Final[TypeAdapter[Path | None]] = TypeAdapter(Path | None)
SHARD_FILES: Final = pytest.StashKey[frozenset[str] | None]()
FAILED_COLLECTIONS: Final = pytest.StashKey[frozenset[str]]()
GROUP_SUFFIX: Final = re.compile(r"@integration_fixture_[0-9a-f]{64}$")
INTEGER_OPTION: Final = TypeAdapter(int)
STRING_ENVIRONMENT: Final = TypeAdapter(dict[str, str])


@runtime_checkable
class WorkerConfig(Protocol):
    @property
    def workeroutput(self) -> dict[str, object]: ...


class WorkerNode(Protocol):
    @property
    def workeroutput(self) -> Mapping[str, object]: ...


def _worker_output(config: object) -> dict[str, object] | None:
    return config.workeroutput if isinstance(config, WorkerConfig) else None


def pytest_addoption(parser: pytest.Parser) -> None:
    parser.addoption("--integration-order-seed", type=int, default=0)
    parser.addoption("--integration-shard-count", type=int, default=1)
    parser.addoption("--integration-shard-index", type=int, default=0)
    parser.addoption("--integration-shard-files", type=Path)


def pytest_configure(config: pytest.Config) -> None:
    config.addinivalue_line("markers", "integration: owned real-service integration contracts")
    config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced")
    config.addinivalue_line("markers", "shares_worker_state: owned proxies here share the worker DB and Redis")
    config.stash[REPORTS] = []
    config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files")))
    config.pluginmanager.register(IntegrationReportPlugin(config))
    build_sdk_models()
    if os.environ.get("INTEGRATION_ROUTING"):
        config.pluginmanager.register(RoutingPlugin(config))


@pytest.hookimpl(optionalhook=True)
def pytest_xdist_setupnodes(config: pytest.Config, specs: Sequence[XSpec]) -> None:
    if (
        os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1"
        or config.getoption("collectonly") is True
        or config.stash[SHARD_FILES] == frozenset()
    ):
        return
    resources: Final = ExitStack()
    config.add_cleanup(resources.close)
    destination: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"])
    stacks: Final = start_worker_stacks(
        tuple(destination / f"gw{index}" for index in range(len(specs))),
        partial(worker_stack, profile=gateway_profile(os.environ)),
        resources,
        os.cpu_count() or 1,
    )
    for spec, stack in zip(specs, stacks, strict=True):
        spec.env = {**STRING_ENVIRONMENT.validate_python(spec.env), **stack}


class IntegrationReportPlugin:
    def __init__(self, config: pytest.Config) -> None:
        self.config = config

    def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
        self.config.stash[REPORTS].append(report)

    def pytest_collectreport(self, report: pytest.CollectReport) -> None:
        if report.failed:
            self.config.stash[FAILED_COLLECTIONS] = self.config.stash.get(FAILED_COLLECTIONS, frozenset()) | {
                report.nodeid.split("::", 1)[0]
            }

    @pytest.hookimpl(optionalhook=True)
    def pytest_xdist_node_collection_finished(self, node: object, ids: Sequence[str]) -> None:
        self.config.stash[COLLECTED] = tuple(_case_identity(nodeid) for nodeid in ids if _owned(nodeid))

    @pytest.hookimpl(optionalhook=True)
    def pytest_testnodedown(self, node: WorkerNode, error: object) -> None:
        if error is not None:
            return
        inventory: Final = WORKER_INVENTORY.validate_python(node.workeroutput.get("integration_inventory"))
        previous: Final = self.config.stash.get(INVENTORY, inventory)
        if inventory != previous:
            raise pytest.UsageError("Integration worker collection inventories disagree")
        self.config.stash[INVENTORY] = inventory

    def pytest_sessionfinish(self, session: pytest.Session, exitstatus: int) -> None:
        output: Final = _worker_output(self.config)
        if output is not None:
            output["integration_inventory"] = self.config.stash.get(INVENTORY, ())


def _case_identity(nodeid: str) -> str:
    return GROUP_SUFFIX.sub("", nodeid)


@pytest.hookimpl(wrapper=True, optionalhook=True)
def pytest_report_from_serializable() -> Generator[
    None, pytest.TestReport | pytest.CollectReport | None, pytest.TestReport | pytest.CollectReport | None
]:
    report: Final = yield
    if not isinstance(report, pytest.TestReport):
        return report
    normalized: Final = copy(report)
    normalized.nodeid = _case_identity(report.nodeid)
    return normalized


@pytest.hookimpl(wrapper=True, optionalhook=True)
def pytest_handlecrashitem(report: pytest.TestReport) -> Generator[None, None, None]:
    result: Final = yield
    report.nodeid = _case_identity(report.nodeid)  # rebind-ok: xdist forwards this same report after the crash hook
    return result


def _owned(nodeid: str) -> bool:
    parts: Final = Path(nodeid.split("::", 1)[0]).parts
    return parts[:2] == ("tests", "integration") and len(parts) > 3 and parts[2] in OWNED_DIRECTORIES


def _digest(seed: int, identity: str) -> bytes:
    return hashlib.sha256(f"{seed}:{identity}".encode()).digest()


def _order_key(seed: int, nodeid: str, group: str = "", inner_group: str = "") -> tuple[bytes, bytes, bytes, bytes]:
    return (
        _digest(seed, nodeid.split("::", 1)[0]),
        _digest(seed, group or nodeid),
        _digest(seed, inner_group or group or nodeid),
        _digest(seed, nodeid),
    )


def _shard_files(path: Path | None) -> frozenset[str] | None:
    if path is None:
        return None
    try:
        files: Final = tuple(path.read_text().splitlines())
    except (OSError, UnicodeError) as error:
        raise pytest.UsageError(f"Cannot read integration shard selection: {error}") from error
    if len(frozenset(files)) != len(files) or any(not file for file in files):
        raise pytest.UsageError("Integration shard selection contains duplicate or blank files")
    return frozenset(files)


@pytest.hookimpl(wrapper=True)
def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> Iterator[None]:
    if config.getoption("dist", default="no") == "loadgroup" or config.getoption("loadgroup", default=False):
        automatic_groups: Final = scheduling_groups(tuple(collected_case(item) for item in items))
        for item in items:
            fixture_group: Final = automatic_groups[item.nodeid]
            if fixture_group and not tuple(item.iter_markers("xdist_group")):
                item.add_marker(
                    pytest.mark.xdist_group("integration_fixture_" + hashlib.sha256(fixture_group.encode()).hexdigest())
                )
    yield
    order_seed: Final = INTEGER_OPTION.validate_python(config.getoption("integration_order_seed"))
    if order_seed:
        cases: Final = tuple(collected_case(item) for item in items)
        groups: Final = scheduling_groups(cases)
        inner_groups: Final = fixture_groups(cases)
        items.sort(
            key=lambda item: _order_key(
                order_seed,
                _case_identity(item.nodeid),
                _case_identity(groups[item.nodeid]),
                _case_identity(inner_groups[item.nodeid]),
            )
        )
    root: Final = Path(__file__).parent
    owned: Final = tuple(
        item
        for item in items
        if item.path.is_relative_to(root) and item.path.relative_to(root).parts[0] in OWNED_DIRECTORIES
    )
    circleci_only: Final = tuple(
        item for item in owned if item.path.relative_to(root.parents[1]).as_posix() not in GITHUB_FILES
    )
    if circleci_only and os.environ.get("GITHUB_ACTIONS") == "true":
        raise pytest.UsageError("Integration contracts are owned by CircleCI")
    for item in owned:
        item.add_marker(pytest.mark.integration)
    config.stash[INVENTORY] = tuple(_case_identity(item.nodeid) for item in owned)
    count: Final = INTEGER_OPTION.validate_python(config.getoption("integration_shard_count"))
    index: Final = INTEGER_OPTION.validate_python(config.getoption("integration_shard_index"))
    if count < 1 or not 0 <= index < count:
        raise pytest.UsageError("Integration shard index must be within the positive shard count")
    files: Final = config.stash[SHARD_FILES]
    if count > 1 and files is None:
        raise pytest.UsageError("Parallel integration runs require CircleCI's shard file selection")
    if files is not None:
        available: Final = frozenset(item.nodeid.split("::", 1)[0] for item in owned)
        for file in sorted(files - available - config.stash.get(FAILED_COLLECTIONS, frozenset())):
            config.hook.pytest_collectreport(
                report=pytest.CollectReport(
                    file, "failed", f"Integration shard selected {file}, which has no tests", []
                )
            )
        selected: Final = tuple(item for item in items if item not in owned or item.nodeid.split("::", 1)[0] in files)
        deselected: Final = tuple(item for item in owned if item.nodeid.split("::", 1)[0] not in files)
        config.hook.pytest_deselected(items=deselected)
        items[:] = selected
    config.stash[COLLECTED] = tuple(_case_identity(item.nodeid) for item in items if item in owned)


def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
    if hasattr(session.config, "workerinput"):
        return
    destination: Final = os.environ.get("INTEGRATION_RESULTS_DIR")
    if destination is None:
        return
    collected: Final = session.config.stash.get(COLLECTED, ())
    reports: Final = tuple(report for report in session.config.stash[REPORTS] if report.nodeid in collected)
    inventory: Final = session.config.stash.get(INVENTORY, collected)
    empty_shard: Final = session.config.stash[SHARD_FILES] == frozenset() and bool(inventory)
    passed: Final = tuple(report.nodeid for report in reports if report.when == "call" and report.passed)
    skipped: Final = tuple(report.nodeid for report in reports if report.skipped)
    complete: Final = (
        (exitstatus == 0 or (empty_shard and exitstatus == pytest.ExitCode.NO_TESTS_COLLECTED))
        and (bool(collected) or empty_shard)
        and sorted(collected) == sorted(passed + skipped)
        and not any(report.failed for report in reports)
    )
    output: Final = Path(destination)
    output.mkdir(parents=True, exist_ok=True)
    (output / "execution.json").write_text(
        json.dumps(
            {
                "collected": collected,
                "inventory": inventory,
                "shard_count": session.config.getoption("integration_shard_count"),
                "shard_index": session.config.getoption("integration_shard_index"),
                "passed": passed,
                "skipped": skipped,
                "complete": complete,
                "exitstatus": 0 if complete else exitstatus,
                "hypothesis_version": version("hypothesis"),
                "hypothesis_seed": session.config.getoption("hypothesis_seed"),
                "order_seed": session.config.getoption("integration_order_seed"),
                "generation": {
                    "max_examples": LIFECYCLE_SETTINGS.max_examples,
                    "stateful_step_count": LIFECYCLE_SETTINGS.stateful_step_count,
                    "database": str(LIFECYCLE_SETTINGS.database),
                    "phases": [phase.name for phase in LIFECYCLE_SETTINGS.phases],
                },
            },
            indent=2,
        )
        + "\n"
    )
    if not complete and exitstatus == 0:
        session.exitstatus = pytest.ExitCode.TESTS_FAILED
    if complete and empty_shard:
        session.exitstatus = pytest.ExitCode.OK


@contextmanager
def _owned_state(item: pytest.Item) -> Iterator[None]:
    with pytest.MonkeyPatch.context() as environment:
        if item.get_closest_marker("shares_worker_state") is not None:
            environment.setenv(OWNED_STATE_VARIABLE, OwnedState.SHARED_WITH_WORKER.value)
        yield


@pytest.hookimpl(wrapper=True)
def pytest_runtest_setup(item: pytest.Item) -> Generator[None, None, None]:
    with _owned_state(item):
        return (yield)


@pytest.hookimpl(wrapper=True)
def pytest_runtest_call(item: pytest.Item) -> Generator[None, None, None]:
    with _owned_state(item):
        return (yield)


@pytest.hookimpl(wrapper=True)
def pytest_runtest_teardown(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, None, None]:
    with _owned_state(item):
        return (yield)


@pytest.fixture
def gateway() -> Iterator[Gateway]:
    with gateway_from_environment() as value:
        httpx.get(f"{value.upstream_url}/__observations", timeout=5, trust_env=False).raise_for_status()
        yield value


@pytest.fixture(scope="session")
def session_gateway() -> Iterator[Gateway]:
    with gateway_from_environment() as value:
        yield value


@pytest.fixture
def scenario(gateway: Gateway) -> Iterator[Scenario]:
    value: Final[Scenario]
    with gateway.scenario() as value:
        yield value


@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def actors(
    session_gateway: Gateway, db: ManagementDatabase
) -> AsyncIterator[ActorSet]:
    with session_gateway.scenario() as scenario:
        from tests.integration._support.actors import create_actor_set

        value: Final[ActorSet] = await create_actor_set(session_gateway, scenario, db)
        yield value


@pytest.fixture(scope="session")
def gateway_client(session_gateway: Gateway, actors: ActorSet) -> GatewayClient:
    from tests.integration._support.actors import create_gateway_client

    return create_gateway_client(session_gateway, actors)


@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def db(session_gateway: Gateway) -> AsyncIterator[ManagementDatabase]:
    from tests.integration._support.actors import management_database

    value: Final[ManagementDatabase]
    async with management_database(session_gateway) as value:
        yield value


@pytest_asyncio.fixture(scope="function", loop_scope="session")
async def resources(db: ManagementDatabase) -> AsyncIterator[ScratchResources]:
    from tests.integration._support.actors import scratch_resources

    value: Final[ScratchResources]
    async with scratch_resources(db) as value:
        yield value


@pytest.fixture(scope="session")
def shared_provider_server() -> Iterator[SharedProvider]:
    if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1":
        pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist")
    with shared_provider() as server:
        yield server
        late: Final = server.received()
        assert late == (), (
            f"the shared fake provider got {[item.target for item in late]} after {server.last_test} finished"
        )


@pytest.fixture
def provider(shared_provider_server: SharedProvider, request: pytest.FixtureRequest) -> Iterator[SharedProvider]:
    stray: Final = shared_provider_server.received()
    shared_provider_server.replies.clear()
    assert stray == (), (
        f"the shared fake provider got {[item.target for item in stray]} "
        f"after {shared_provider_server.last_test} finished"
    )
    yield shared_provider_server
    shared_provider_server.last_test = request.node.nodeid
    unused: Final = len(shared_provider_server.replies)
    unread: Final = shared_provider_server.received()
    shared_provider_server.replies.clear()
    assert unused == 0, f"{unused} queued provider replies were never requested"
    assert unread == (), f"the test never read the provider requests {[item.target for item in unread]}"


@pytest.fixture
def peer(gateway: Gateway) -> Iterator[Gateway]:
    url: Final = os.environ["INTEGRATION_PEER_URL"]
    assert url.rstrip("/") != str(gateway.client.base_url).rstrip("/")
    with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
        eventually(lambda: cache.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1], lambda count: count >= 2)
    with httpx.Client(base_url=url, timeout=15, trust_env=False) as client:
        yield Gateway(client, gateway.key, gateway.upstream_url)