comet-ml--opik/tests_load/suite/python_sdk/_helpers.py

"""Helpers shared by the load-test suite.

Functions here are intentionally simple so individual tests can be read
top-to-bottom without jumping around.
"""

import base64
import contextlib
import json
import logging
import os
import random
import string
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, Iterator, List, Set, Tuple

import opik
from opik import Opik
from opik.rest_api.types.span_public import SpanPublic
from opik.rest_api.types.trace_public import TracePublic


PROJECT_PREFIX: str = "loadtest"
REPORT_DIR: Path = Path(__file__).resolve().parents[2] / ".last_run"

KB: int = 1_000
MB: int = 1_000_000

REQUIRED_SPAN_FIELDS: Tuple[str, ...] = ("name", "end_time", "input", "output")


class Metrics:
    """Tiny key/value recorder with a `timer(label)` context manager.

    One instance per test; written to ``REPORT_DIR/<test_name>.json`` at
    teardown by the ``metrics`` fixture.
    """

    test_name: str
    _data: Dict[str, Any]

    def __init__(self, test_name: str) -> None:
        self.test_name = test_name
        self._data = {"test_name": test_name}

    def __setitem__(self, key: str, value: Any) -> None:
        self._data[key] = value

    @contextlib.contextmanager
    def timer(self, label: str) -> Iterator[None]:
        start: float = time.perf_counter()
        yield
        self._data[f"{label}_seconds"] = round(time.perf_counter() - start, 3)

    def write(self) -> None:
        REPORT_DIR.mkdir(parents=True, exist_ok=True)
        out_path: Path = REPORT_DIR / f"{self.test_name}.json"
        out_path.write_text(json.dumps(self._data, indent=2, default=str))
        logging.getLogger(self.test_name).info(
            "metrics: %s", json.dumps(self._data, default=str)
        )


def unique_project_name(scenario: str) -> str:
    timestamp: str = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
    return f"{PROJECT_PREFIX}-{scenario}-{timestamp}-{uuid.uuid4().hex[:6]}"


def random_text(size_bytes: int) -> str:
    return "".join(random.choices(string.ascii_letters + string.digits, k=size_bytes))


def random_bytes(size_bytes: int) -> bytes:
    return os.urandom(size_bytes)


PNG_MAGIC: bytes = b"\x89PNG\r\n\x1a\n"


def random_base64_png(size_bytes: int) -> str:
    """Returns a standard-alphabet base64-encoded byte stream that the SDK's
    attachment-extraction pipeline recognises as a PNG.

    Two constraints from the SDK that the naïve approach (random bytes,
    ``urlsafe_b64encode``) trips over:

    1. The extraction regex is ``[A-Za-z0-9+/]``, so the encoded blob
       must use the standard ``+/`` alphabet (``b64encode``), not the
       ``-_`` url-safe one.
    2. The decoder MIME-sniffs the raw bytes and skips anything that
       comes back as ``application/octet-stream`` or ``text/plain``,
       so the byte stream must start with a recognised magic header.

    The bytes after the 8-byte PNG magic are pure noise — the test
    doesn't render the image, it just exercises the extract/upload path.
    """
    raw_size: int = (size_bytes // 4) * 3
    payload: bytes = PNG_MAGIC + os.urandom(max(0, raw_size - len(PNG_MAGIC)))
    return base64.b64encode(payload).decode("ascii")


def think_time() -> None:
    """Sleeps a small randomized interval so submits aren't lockstep.

    Real callers always have at least some gap between consecutive SDK calls
    (LLM responses, request handling, etc.). The range is small on purpose so
    it doesn't dominate runtime for high-count scenarios.
    """
    time.sleep(random.uniform(0.0005, 0.002))


def opik_client() -> Opik:
    return opik.Opik()


def now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _assert_required_fields_set(
    items: Iterable[Any], required_fields: Tuple[str, ...], kind: str
) -> None:
    """Fails the test if any item is missing one of `required_fields`.

    Catches partial-write bugs where the entity lands but a later update
    (e.g. ``end_time``) was dropped or never sent.
    """
    broken: List[str] = []
    for item in items:
        for field in required_fields:
            if getattr(item, field) is None:
                broken.append(f"{item.id}: {field} is None")
                break
    if broken:
        raise AssertionError(
            f"{len(broken)} {kind}(s) missing required fields "
            f"({', '.join(required_fields)}); sample: {broken[:5]}"
        )


def verify_traces(
    client: Opik,
    project_name: str,
    expected_count: int,
    timeout_seconds: int = 900,
) -> List[TracePublic]:
    """Polls until at least `expected_count` traces are visible in the project.

    Excludes ``input``/``output``/``metadata`` from the response. We only
    need ``id``, ``name``, ``end_time`` for the load-test verification,
    and excluding the bulky fields:

    1. Streams much less data back for the high-count scenarios.
    2. Side-steps OPIK-6651, where attachment-extracted traces fail to
       stream because ``AttachmentService.list`` can't read
       ``workspaceName`` from the reactor context during enrichment.
       Excluding the fields skips the enrichment path entirely. Drop
       the exclude list once OPIK-6651 is fixed.

    ``end_time``/``name`` field checks still run; ``input``/``output``
    are not included in the response so we can't assert on them here.
    """
    traces: List[TracePublic] = client.search_traces(
        project_name=project_name,
        max_results=expected_count,
        wait_for_at_least=expected_count,
        wait_for_timeout=timeout_seconds,
        exclude=["input", "output", "metadata"],
    )
    _assert_required_fields_set(traces, ("name", "end_time"), "trace")
    return traces


def verify_spans_for_trace(
    client: Opik,
    project_name: str,
    trace_id: str,
    expected_count: int,
    timeout_seconds: int = 600,
) -> List[SpanPublic]:
    """Polls until at least `expected_count` spans are visible for a trace."""
    spans: List[SpanPublic] = client.search_spans(
        project_name=project_name,
        trace_id=trace_id,
        max_results=expected_count,
        wait_for_at_least=expected_count,
        wait_for_timeout=timeout_seconds,
    )
    _assert_required_fields_set(spans, REQUIRED_SPAN_FIELDS, "span")
    return spans


def verify_exact_trace_ids(
    client: Opik,
    project_name: str,
    expected_ids: Set[str],
    timeout_seconds: int = 900,
) -> Set[str]:
    """Waits for all `expected_ids` to land and returns the delivered id set.

    Raises ``AssertionError`` listing a sample of missing ids if some traces
    are never delivered within the timeout. This catches dropped messages
    (e.g. regression cases like OPIK-6444) rather than just under-counts.

    Required-field validation runs inside ``verify_traces`` so any trace
    landing without ``end_time``, ``name``, ``input``, or ``output`` also
    fails the test.
    """
    traces: List[TracePublic] = verify_traces(
        client,
        project_name=project_name,
        expected_count=len(expected_ids),
        timeout_seconds=timeout_seconds,
    )
    delivered: Set[str] = {trace.id for trace in traces}
    missing: Set[str] = expected_ids - delivered
    if missing:
        raise AssertionError(
            f"{len(missing)}/{len(expected_ids)} traces missing after flush; "
            f"sample missing ids: {sorted(missing)[:3]}"
        )
    return delivered


def verify_attachments(
    client: Opik,
    project_name: str,
    entity_type: str,
    entity_id: str,
    expected_count: int,
    timeout_seconds: int = 600,
) -> int:
    """Polls the attachments REST endpoint until `expected_count` are listed.

    The endpoint's ``path`` query parameter is a base64-encoded base URL
    the backend uses to build download links — we pass the SDK's
    configured base URL so the round-trip succeeds. Uses the standard
    ``b64encode`` alphabet (not ``urlsafe_b64encode``) to match the
    contract used by ``attachment/client.py`` and ``tests/e2e/verifiers.py``.
    """
    project = client.rest_client.projects.retrieve_project(name=project_name)
    base_url: str = str(client.config.url_override).rstrip("/")
    encoded_base_url: str = base64.b64encode(
        base_url.encode("utf-8")
    ).decode("ascii")
    deadline: float = time.time() + timeout_seconds
    last_seen: int = 0
    while time.time() < deadline:
        page = client.rest_client.attachments.attachment_list(
            project_id=project.id,
            entity_type=entity_type,
            entity_id=entity_id,
            path=encoded_base_url,
            size=expected_count + 10,
        )
        last_seen = len(page.content or [])
        if last_seen >= expected_count:
            return last_seen
        time.sleep(1)
    raise TimeoutError(
        f"Only {last_seen}/{expected_count} attachments visible after "
        f"{timeout_seconds}s for {entity_type} {entity_id}"
    )