chore: import upstream snapshot with attribution

sdks/python/tests/testlib/__init__.py

@@ -0,0 +1,32 @@
+from .any_compare_helpers import ANY, ANY_BUT_NONE, ANY_DICT, ANY_LIST, ANY_STRING
+from .assert_helpers import (
+    assert_dict_has_keys,
+    assert_dict_keys_in_list,
+    assert_dicts_equal,
+    assert_equal,
+)
+from .backend_emulator_message_processor import BackendEmulatorMessageProcessor
+from .concurrency_helpers import ThreadSafeCounter
+from .project_naming import generate_project_name
+from .models import AttachmentModel, FeedbackScoreModel, SpanModel, TraceModel
+from .patch_helpers import patch_environ
+
+__all__ = [
+    "ANY",
+    "ANY_BUT_NONE",
+    "ANY_DICT",
+    "ANY_LIST",
+    "ANY_STRING",
+    "AttachmentModel",
+    "BackendEmulatorMessageProcessor",
+    "ThreadSafeCounter",
+    "FeedbackScoreModel",
+    "SpanModel",
+    "TraceModel",
+    "assert_dict_has_keys",
+    "assert_dict_keys_in_list",
+    "assert_dicts_equal",
+    "assert_equal",
+    "generate_project_name",
+    "patch_environ",
+]

sdks/python/tests/testlib/any_compare_helpers.py

@@ -0,0 +1,115 @@
+from typing import Optional, Dict
+
+from unittest import mock
+
+
+class AnyButNone:
+    """A helper object that compares equal to everything but None."""
+
+    def __eq__(self, other):
+        if other is None:
+            return False
+
+        return True
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        return "<ANY_BUT_NONE>"
+
+
+class AnyDict:
+    """A helper object that compares equal to all dicts."""
+
+    def __init__(self, containing: Optional[Dict] = None):
+        self.containing_items = containing
+
+    def __eq__(self, other):
+        if not isinstance(other, dict):
+            return False
+
+        if self.containing_items is None:
+            return True
+
+        if other.items() >= self.containing_items.items():
+            return True
+
+        return False
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        if self.containing_items is None:
+            return "<ANY_DICT>"
+        return "<ANY_DICT_WITH_CONTAIN_CONDITION>"
+
+    def containing(self, containing: Dict):
+        return AnyDict(containing)
+
+
+class AnyList:
+    """A helper object that compares equal to all lists."""
+
+    def __eq__(self, other):
+        if isinstance(other, list):
+            return True
+
+        return False
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        return "<ANY_LIST>"
+
+
+class AnyString:
+    """A helper object that provides partial equality check to strings."""
+
+    def __init__(
+        self, startswith: Optional[str] = None, containing: Optional[str] = None
+    ):
+        self._startswith = startswith
+        self._containing = containing
+
+    def __eq__(self, other):
+        if not isinstance(other, str):
+            return False
+
+        if self._startswith is None and self._containing is None:
+            return True
+
+        if self._startswith is not None and not other.startswith(self._startswith):
+            return False
+
+        if self._containing is not None and self._containing not in other:
+            return False
+
+        return True
+
+    def __repr__(self):
+        conditions = []
+        if self._startswith is not None:
+            conditions.append(f"startswith='{self._startswith}'")
+        if self._containing is not None:
+            conditions.append(f"containing='{self._containing}'")
+
+        if not conditions:
+            return "<ANY_STRING>"
+
+        return f"<ANY_STRING({', '.join(conditions)})>"
+
+    def starting_with(self, startswith: str):
+        return AnyString(startswith=startswith, containing=self._containing)
+
+    def containing(self, containing: str):
+        return AnyString(startswith=self._startswith, containing=containing)
+
+
+ANY = mock.ANY
+ANY_BUT_NONE = AnyButNone()
+ANY_DICT = AnyDict()
+ANY_LIST = AnyList()
+ANY_STRING = AnyString()

sdks/python/tests/testlib/assert_helpers.py

@@ -0,0 +1,75 @@
+from typing import List, Any, Optional, Dict, Mapping
+
+import logging
+import pytest_deepassert
+
+from opik.evaluation.metrics import score_result
+
+LOGGER = logging.getLogger(__name__)
+
+
+def assert_equal(expected: Any, actual: Any) -> None:
+    __tracebackhide__ = True
+
+    pytest_deepassert.equal(expected, actual)
+
+
+def assert_dicts_equal(
+    dict1: Mapping[str, Any],
+    dict2: Mapping[str, Any],
+    ignore_keys: Optional[List[str]] = None,
+) -> None:
+    __tracebackhide__ = True
+
+    dict1_copy, dict2_copy = {**dict1}, {**dict2}
+
+    ignore_keys = [] if ignore_keys is None else ignore_keys
+
+    for key in ignore_keys:
+        dict1_copy.pop(key, None)
+        dict2_copy.pop(key, None)
+
+    pytest_deepassert.equal(dict1_copy, dict2_copy)
+
+
+def assert_dict_has_keys(dic: Dict[str, Any], keys: List[str]) -> None:
+    dict_has_keys = all(key in dic for key in keys)
+
+    if dict_has_keys:
+        return
+
+    raise AssertionError(
+        f"Dict doesn't contain all the required keys. Dict keys: {dic.keys()}, required keys: {keys}"
+    )
+
+
+def assert_dict_keys_in_list(dic: Dict[str, Any], keys: List[str]) -> None:
+    """
+    Asserts that all keys in the dictionary are present in the given list.
+
+    Args:
+        dic: The dictionary whose keys need to be checked
+        keys: The list of allowed keys
+
+    Raises:
+        AssertionError: If any key in the dictionary is not in the provided list
+    """
+    invalid_keys = [key for key in dic.keys() if key not in keys]
+
+    if len(invalid_keys) == 0:
+        return
+
+    raise AssertionError(
+        f"Dict contains keys that are not in the allowed list. Invalid keys: {invalid_keys}, allowed keys: {keys}"
+    )
+
+
+def assert_score_result(
+    result: score_result.ScoreResult, include_reason: bool = True
+) -> None:
+    assert result.scoring_failed is False
+    assert isinstance(result.value, float)
+    assert 0.0 <= result.value <= 1.0
+    if include_reason:
+        assert isinstance(result.reason, str)
+        assert len(result.reason) > 0

sdks/python/tests/testlib/backend_emulator_message_processor.py

@@ -0,0 +1,199 @@
+import datetime
+from typing import List, Dict, Optional, Any, TYPE_CHECKING
+
+from opik.message_processing.emulation import emulator_message_processor
+from opik.types import ErrorInfoDict, SpanType, TraceSource
+from . import models
+
+if TYPE_CHECKING:
+    from . import noop_file_upload_manager
+
+
+class BackendEmulatorMessageProcessor(
+    emulator_message_processor.EmulatorMessageProcessor
+):
+    """
+    This class serves as a replacement for the real backend. It collects all logged messages
+    to be used in tests.
+
+    Optionally accepts a file_upload_manager to access attachment data that was
+    intercepted by the FileUploadPreprocessor before reaching the message processor.
+    """
+
+    def __init__(
+        self,
+        active: bool = True,
+        merge_duplicates: bool = True,
+        file_upload_manager: Optional[
+            "noop_file_upload_manager.FileUploadManagerEmulator"
+        ] = None,
+    ) -> None:
+        super().__init__(active=active, merge_duplicates=merge_duplicates)
+        self._file_upload_manager = file_upload_manager
+
+    def create_trace_model(
+        self,
+        trace_id: str,
+        start_time: datetime.datetime,
+        name: Optional[str],
+        project_name: str,
+        input: Any,
+        output: Any,
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        end_time: Optional[datetime.datetime],
+        spans: Optional[List[models.SpanModel]],
+        feedback_scores: Optional[List[models.FeedbackScoreModel]],
+        error_info: Optional[ErrorInfoDict],
+        thread_id: Optional[str],
+        source: TraceSource,
+        last_updated_at: Optional[datetime.datetime] = None,
+        environment: Optional[str] = None,
+    ) -> models.TraceModel:
+        if spans is None:
+            spans = []
+        if feedback_scores is None:
+            feedback_scores = []
+
+        return models.TraceModel(
+            id=trace_id,
+            start_time=start_time,
+            name=name,
+            project_name=project_name,
+            input=input,
+            output=output,
+            tags=tags,
+            metadata=metadata,
+            end_time=end_time,
+            spans=spans,
+            feedback_scores=feedback_scores,
+            error_info=error_info,
+            thread_id=thread_id,
+            last_updated_at=last_updated_at,
+            source=source,
+            environment=environment,
+        )
+
+    def create_span_model(
+        self,
+        span_id: str,
+        start_time: datetime.datetime,
+        name: Optional[str],
+        input: Any,
+        output: Any,
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        type: SpanType,
+        usage: Optional[Dict[str, Any]],
+        end_time: Optional[datetime.datetime],
+        project_name: str,
+        spans: Optional[List[models.SpanModel]],
+        feedback_scores: Optional[List[models.FeedbackScoreModel]],
+        model: Optional[str],
+        provider: Optional[str],
+        error_info: Optional[ErrorInfoDict],
+        total_cost: Optional[float],
+        last_updated_at: Optional[datetime.datetime],
+        source: TraceSource,
+        environment: Optional[str] = None,
+    ) -> models.SpanModel:
+        if spans is None:
+            spans = []
+        if feedback_scores is None:
+            feedback_scores = []
+
+        return models.SpanModel(
+            id=span_id,
+            start_time=start_time,
+            name=name,
+            input=input,
+            output=output,
+            tags=tags,
+            metadata=metadata,
+            type=type,
+            usage=usage,
+            end_time=end_time,
+            project_name=project_name,
+            spans=spans,
+            feedback_scores=feedback_scores,
+            model=model,
+            provider=provider,
+            error_info=error_info,
+            total_cost=total_cost,
+            last_updated_at=last_updated_at,
+            source=source,
+            environment=environment,
+        )
+
+    def create_feedback_score_model(
+        self,
+        score_id: str,
+        name: str,
+        value: float,
+        category_name: Optional[str],
+        reason: Optional[str],
+    ) -> models.FeedbackScoreModel:
+        return models.FeedbackScoreModel(
+            id=score_id,
+            name=name,
+            value=value,
+            category_name=category_name,
+            reason=reason,
+        )
+
+    @property
+    def trace_trees(self) -> List[models.TraceModel]:
+        """
+        Override to add attachments from the file upload manager.
+
+        Attachments are intercepted by FileUploadPreprocessor before reaching
+        the message processor, so we need to get them from the upload manager.
+        """
+        # Get base trace trees from parent
+        traces = super().trace_trees
+
+        # If we have a file upload manager, add attachments to spans and traces
+        if self._file_upload_manager is not None:
+            self._add_attachments_to_traces(traces)
+
+        return traces
+
+    def _add_attachments_to_traces(self, traces: List[models.TraceModel]) -> None:
+        """Add attachments from file upload manager to traces and their spans."""
+        for trace in traces:
+            # Add trace-level attachments
+            trace_attachments = self._file_upload_manager.attachments_by_trace.get(
+                trace.id, []
+            )
+            if trace_attachments:
+                trace.attachments = [
+                    models.AttachmentModel(
+                        file_path=att.file_path,
+                        file_name=att.file_name,
+                        content_type=att.mime_type or "",
+                    )
+                    for att in trace_attachments
+                ]
+
+            # Add span-level attachments recursively
+            self._add_attachments_to_spans(trace.spans)
+
+    def _add_attachments_to_spans(self, spans: List[models.SpanModel]) -> None:
+        """Recursively add attachments to spans."""
+        for span in spans:
+            span_attachments = self._file_upload_manager.attachments_by_span.get(
+                span.id, []
+            )
+            if span_attachments:
+                span.attachments = [
+                    models.AttachmentModel(
+                        file_path=att.file_path,
+                        file_name=att.file_name,
+                        content_type=att.mime_type or "",
+                    )
+                    for att in span_attachments
+                ]
+
+            # Recurse into nested spans
+            if span.spans:
+                self._add_attachments_to_spans(span.spans)

sdks/python/tests/testlib/concurrency_helpers.py

@@ -0,0 +1,20 @@
+import threading
+
+
+class ThreadSafeCounter:
+    """Thread-safe counter for tracking invocations in concurrent tests."""
+
+    def __init__(self) -> None:
+        self._value = 0
+        self._lock = threading.Lock()
+
+    def increment(self) -> int:
+        """Increment and return the new value (1-based)."""
+        with self._lock:
+            self._value += 1
+            return self._value
+
+    @property
+    def value(self) -> int:
+        with self._lock:
+            return self._value

sdks/python/tests/testlib/environment.py

@@ -0,0 +1,19 @@
+import os
+
+
+def has_openai_api_key():
+    """
+    Checks if the OpenAI API key and organization ID exist in the environment variables.
+
+    This function verifies the presence and non-emptiness of both the 'OPENAI_API_KEY'
+    and 'OPENAI_ORG_ID' environment variables in the system's environment.
+
+    Returns:
+        bool: True if both variables exist and are not empty, False otherwise.
+    """
+    return (
+        "OPENAI_API_KEY" in os.environ
+        and len(os.environ["OPENAI_API_KEY"]) > 0
+        and "OPENAI_ORG_ID" in os.environ
+        and len(os.environ["OPENAI_ORG_ID"]) > 0
+    )

sdks/python/tests/testlib/fake_message_factory.py

@@ -0,0 +1,270 @@
+import dataclasses
+import random
+import uuid
+from datetime import datetime, timedelta
+from typing import List, Optional
+
+from opik.message_processing import messages
+from opik.types import ErrorInfoDict
+
+
+@dataclasses.dataclass
+class LongStr:
+    value: str
+
+    def __str__(self) -> str:
+        return self.value[1] + ".." + self.value[-1]
+
+    def __repr__(self) -> str:
+        return str(self)
+
+
+ONE_KILOBYTE = 1024
+ONE_MEGABYTE = ONE_KILOBYTE * ONE_KILOBYTE
+
+
+def fake_create_trace_message_batch(
+    count: int = 1000,
+    approximate_trace_size: int = ONE_MEGABYTE,
+    has_ended: Optional[bool] = None,
+) -> List[messages.CreateTraceMessage]:
+    """
+    Factory method to create a batch with a specified number of
+    CreateTraceMessage objects initialized with fake data.
+
+    Args:
+        approximate_trace_size: The approximate size of each trace in megabytes
+        count: Number of CreateTraceMessage objects to include in the batch (default: 1000)
+        has_ended: the flag to indicate if the trace has ended. If None, the trace will
+            be randomly decided to be ended or not.
+
+    Returns:
+        CreateTraceBatchMessage containing the specified number of fake CreateTraceMessage objects
+    """
+    dummy_traces = []
+
+    for i in range(count):
+        # Generate a unique trace ID
+        trace_id = str(uuid.uuid4())
+
+        # Create a random start time within the last 24 hours
+        start_time = datetime.now() - timedelta(
+            hours=random.randint(0, 23),
+            minutes=random.randint(0, 59),
+            seconds=random.randint(0, 59),
+        )
+
+        # Randomly decide if the trace has ended
+        if has_ended is None:
+            has_ended = random.choice([True, False])
+
+        if has_ended:
+            end_time = start_time + timedelta(seconds=random.randint(1, 3600))
+            last_updated_at = end_time
+        else:
+            end_time = None
+            last_updated_at = start_time
+
+        # Generate dummy input data
+        input_data = {
+            "prompt": f"This is a dummy prompt #{i}",
+            "parameters": {
+                "temperature": round(random.uniform(0.1, 1.0), 2),
+                "max_tokens": random.randint(10, 1000),
+                "long_string": LongStr("a" * approximate_trace_size),
+            },
+        }
+
+        # Generate dummy output data if the trace has ended
+        output_data = (
+            {
+                "response": f"This is a dummy response for prompt #{i}",
+                "tokens_used": random.randint(10, 500),
+            }
+            if has_ended
+            else None
+        )
+
+        # Generate dummy metadata
+        metadata = {
+            "model": random.choice(["gpt-3.5-turbo", "gpt-4", "claude-2", "llama-2"]),
+            "environment": random.choice(["production", "staging", "development"]),
+            "client_id": f"client-{random.randint(1000, 9999)}",
+        }
+
+        # Generate random tags
+        available_tags = [
+            "important",
+            "experiment",
+            "production",
+            "test",
+            "debug",
+            "high-priority",
+            "low-priority",
+        ]
+        tags = random.sample(
+            available_tags, k=random.randint(0, min(3, len(available_tags)))
+        )
+
+        # Randomly decide if there's an error
+        has_error = random.random() < 0.1  # 10% chance of error
+        error_info = (
+            ErrorInfoDict(
+                exception_type=random.choice(
+                    ["TimeoutError", "ValidationError", "AuthenticationError"]
+                ),
+                traceback=f"Dummy stacktrace for error in trace #{i}",
+            )
+            if has_error
+            else None
+        )
+
+        # Generate a thread ID for some traces
+        thread_id = (
+            str(uuid.uuid4()) if random.random() < 0.7 else None
+        )  # 70% chance of having a thread ID
+
+        # Create the trace message
+        trace_message = messages.CreateTraceMessage(
+            trace_id=trace_id,
+            project_name="dummy-project",
+            name=f"Dummy Trace #{i}",
+            start_time=start_time,
+            end_time=end_time,
+            input=input_data,
+            output=output_data,
+            metadata=metadata,
+            tags=tags,
+            error_info=error_info,
+            thread_id=thread_id,
+            last_updated_at=last_updated_at,
+            source="sdk",
+        )
+
+        dummy_traces.append(trace_message)
+
+    return dummy_traces
+
+
+def fake_span_create_message_batch(
+    count: int = 1000,
+    approximate_span_size: int = ONE_MEGABYTE,
+    has_ended: Optional[bool] = None,
+) -> List[messages.CreateSpanMessage]:
+    """
+    Factory method to create a list with a specified number of
+    CreateSpanMessage objects initialized with fake data.
+
+    Args:
+        approximate_span_size: The approximate size of each span in megabytes
+        count: Number of CreateSpanMessage objects to include in the batch (default: 1000)
+        has_ended: the flag to indicate if the span has ended. If None, the span will
+            be randomly decided to be ended or not.
+
+    Returns:
+        CreateSpansBatchMessage containing the specified number of fake CreateSpanMessage objects
+    """
+    dummy_spans = []
+
+    for i in range(count):
+        # Generate a unique span ID
+        span_id = str(uuid.uuid4())
+
+        # Create a random start time within the last 24 hours
+        start_time = datetime.now() - timedelta(
+            hours=random.randint(0, 23),
+            minutes=random.randint(0, 59),
+            seconds=random.randint(0, 59),
+        )
+
+        # Randomly decide if the span has ended
+        if has_ended is None:
+            has_ended = random.choice([True, False])
+
+        if has_ended:
+            end_time = start_time + timedelta(seconds=random.randint(1, 3600))
+            last_updated_at = end_time
+        else:
+            end_time = None
+            last_updated_at = start_time
+
+        # Generate dummy input data
+        input_data = {
+            "prompt": f"This is a dummy prompt #{i}",
+            "parameters": {
+                "temperature": round(random.uniform(0.1, 1.0), 2),
+                "max_tokens": random.randint(10, 1000),
+                "long_string": LongStr("a" * approximate_span_size),
+            },
+        }
+
+        # Generate dummy output data if the span has ended
+        output_data = (
+            {
+                "response": f"This is a dummy response for prompt #{i}",
+                "tokens_used": random.randint(10, 500),
+            }
+            if has_ended
+            else None
+        )
+
+        # Generate dummy metadata
+        metadata = {
+            "model": random.choice(["gpt-3.5-turbo", "gpt-4", "claude-2", "llama-2"]),
+            "environment": random.choice(["production", "staging", "development"]),
+            "client_id": f"client-{random.randint(1000, 9999)}",
+        }
+
+        # Generate random tags
+        available_tags = [
+            "important",
+            "experiment",
+            "production",
+            "test",
+            "debug",
+            "high-priority",
+            "low-priority",
+        ]
+        tags = random.sample(
+            available_tags, k=random.randint(0, min(3, len(available_tags)))
+        )
+
+        # Randomly decide if there's an error
+        has_error = random.random() < 0.1  # 10% chance of error
+        error_info = (
+            ErrorInfoDict(
+                exception_type=random.choice(
+                    ["TimeoutError", "ValidationError", "AuthenticationError"]
+                ),
+                traceback=f"Dummy stacktrace for error in trace #{i}",
+            )
+            if has_error
+            else None
+        )
+
+        # Create the span message
+        span_message = messages.CreateSpanMessage(
+            span_id=span_id,
+            trace_id=str(uuid.uuid4()),
+            parent_span_id=span_id,  # This is wrong, but it's okay for dummy data
+            project_name="dummy-project",
+            name=f"Dummy Span #{i}",
+            start_time=start_time,
+            end_time=end_time,
+            input=input_data,
+            output=output_data,
+            metadata=metadata,
+            tags=tags,
+            error_info=error_info,
+            type="general",
+            usage=None,
+            model=metadata["model"],
+            provider=None,
+            total_cost=random.random() * 0.01,
+            last_updated_at=last_updated_at,
+            source="sdk",
+        )
+
+        dummy_spans.append(span_message)
+
+    return dummy_spans

sdks/python/tests/testlib/models.py

@@ -0,0 +1,42 @@
+import dataclasses
+import datetime
+from typing import Optional, List
+
+from opik.message_processing.emulation import models
+from opik.types import TraceSource
+from .any_compare_helpers import ANY, ANY_BUT_NONE
+
+
+@dataclasses.dataclass
+class SpanModel(models.SpanModel):
+    project_name: str = dataclasses.field(
+        default_factory=lambda: ANY
+    )  # we don't want to check the project name unless it's specified explicitly in the test
+    last_updated_at: Optional[datetime.datetime] = dataclasses.field(
+        default_factory=lambda: ANY_BUT_NONE
+    )  # we don't want to check the last_updated_at unless it's specified explicitly in the test - just make sure it's not None
+    attachments: Optional[List[models.AttachmentModel]] = dataclasses.field(
+        default_factory=lambda: ANY
+    )  # we don't want to check attachments unless explicitly specified in the test
+    source: TraceSource = dataclasses.field(default_factory=lambda: "sdk")
+
+
+@dataclasses.dataclass
+class TraceModel(models.TraceModel):
+    project_name: str = dataclasses.field(
+        default_factory=lambda: ANY
+    )  # we don't want to check the project name unless it's specified explicitly in the test
+    attachments: Optional[List[models.AttachmentModel]] = dataclasses.field(
+        default_factory=lambda: ANY
+    )  # we don't want to check attachments unless explicitly specified in the test
+    source: TraceSource = dataclasses.field(default_factory=lambda: "sdk")
+
+
+@dataclasses.dataclass
+class FeedbackScoreModel(models.FeedbackScoreModel):
+    pass
+
+
+@dataclasses.dataclass
+class AttachmentModel(models.AttachmentModel):
+    pass

sdks/python/tests/testlib/noop_file_upload_manager.py

@@ -0,0 +1,63 @@
+from typing import Dict, List, Optional
+
+from opik.file_upload import base_upload_manager, types as upload_types
+from opik.message_processing import messages
+
+
+class FileUploadManagerEmulator(base_upload_manager.BaseFileUploadManager):
+    """
+    A file upload manager emulator that stores attachment messages in memory.
+
+    Attachments are stored by entity_id (span_id or trace_id) for easy lookup
+    by the test backend emulator.
+    """
+
+    def __init__(self) -> None:
+        self.current_uploads: List[messages.BaseMessage] = []
+        # Store attachments by entity_id for lookup
+        self.attachments_by_span: Dict[str, List[messages.CreateAttachmentMessage]] = {}
+        self.attachments_by_trace: Dict[
+            str, List[messages.CreateAttachmentMessage]
+        ] = {}
+
+    def upload(
+        self,
+        message: messages.BaseMessage,
+        on_upload_success: Optional[upload_types.OnUploadSuccessCallback],
+        on_upload_failed: Optional[upload_types.OnUploadFailureCallback],
+    ) -> None:
+        self.current_uploads.append(message)
+
+        # Store attachment messages by entity for easy lookup
+        if isinstance(message, messages.CreateAttachmentMessage):
+            if message.entity_type == "span":
+                if message.entity_id not in self.attachments_by_span:
+                    self.attachments_by_span[message.entity_id] = []
+                self.attachments_by_span[message.entity_id].append(message)
+            elif message.entity_type == "trace":
+                if message.entity_id not in self.attachments_by_trace:
+                    self.attachments_by_trace[message.entity_id] = []
+                self.attachments_by_trace[message.entity_id].append(message)
+
+            if on_upload_success is not None:
+                on_upload_success()
+
+    def remaining_data(self) -> base_upload_manager.RemainingUploadData:
+        return base_upload_manager.RemainingUploadData(
+            uploads=len(self.current_uploads), bytes=-1, total_size=-1
+        )
+
+    def flush(self, timeout: Optional[float], sleep_time: int = 5) -> bool:
+        self.current_uploads = []
+        return True
+
+    def all_done(self) -> bool:
+        return len(self.current_uploads) == 0
+
+    def failed_uploads(self, timeout: Optional[float]) -> int:
+        return 0
+
+    def close(self) -> None:
+        self.current_uploads = []
+        self.attachments_by_span = {}
+        self.attachments_by_trace = {}

sdks/python/tests/testlib/patch_helpers.py

@@ -0,0 +1,31 @@
+import contextlib
+import os
+from typing import Any, Dict, List
+
+
+@contextlib.contextmanager
+def patch_environ(
+    add_keys: Dict[str, Any],
+    remove_keys: List[str] = None,
+):
+    """
+    Temporarily set environment variables inside the context manager and
+    fully restore the previous environment afterward
+    """
+    original_env = {key: os.getenv(key) for key in add_keys}
+
+    for key in remove_keys or []:
+        if key in os.environ:
+            original_env[key] = os.getenv(key)
+            del os.environ[key]
+
+    os.environ.update(add_keys)
+
+    try:
+        yield
+    finally:
+        for key, value in original_env.items():
+            if value is None:
+                del os.environ[key]
+            else:
+                os.environ[key] = value

sdks/python/tests/testlib/project_naming.py

@@ -0,0 +1,49 @@
+"""Resource-name helper for the e2e test suite.
+
+Project names are composed as ``prefix1-prefix2-...-random_chars`` so
+re-running the suite against the same backend never reuses a name.
+
+NOTE on pytest-xdist:
+``random_chars`` is non-deterministic, so two xdist workers compute
+different names for the same call. That is fine for module-level
+constants used only inside test bodies (under ``--dist=loadfile`` each
+file runs on a single worker, so other workers' copies of the constant
+are never used). It is *not* fine for constants embedded in
+``@pytest.mark.parametrize`` values: every worker collects every
+parametrize id, and xdist refuses to run when ids differ across
+workers. Use a plain string literal there instead.
+"""
+
+import secrets
+
+
+def _random_chars(n: int = 6) -> str:
+    # Duplicates tests/conftest.py::random_chars on purpose: importing
+    # from tests/conftest.py here would create a cycle (conftest itself
+    # imports from testlib). `secrets` (vs `random`) keeps per-worker
+    # uniqueness independent of any seeding in the test environment
+    # (e.g. a stray random.seed(...) in a fixture, or a pinned
+    # PYTHONHASHSEED) — load-bearing for the xdist isolation contract.
+    return secrets.token_hex((n + 1) // 2)[:n]
+
+
+def generate_project_name(*prefixes: str) -> str:
+    """Return ``prefix1-prefix2-...-random_chars``.
+
+    Use at module top in any e2e test file that needs to reference its
+    project — usually paired with ``__name__`` so the name embeds the
+    test file::
+
+        PROJECT_NAME = generate_project_name("e2e", __name__)
+
+    The e2e ``configure_e2e_tests_env`` fixture reads ``PROJECT_NAME``
+    from each test module and patches ``OPIK_PROJECT_NAME`` to that
+    value, so the constant is the single source of truth.
+
+    Dotted segments (e.g. ``__name__`` is ``tests.e2e.test_dataset``)
+    are reduced to their last component so the resulting name stays a
+    single hyphenated string.
+    """
+    parts = [prefix.rsplit(".", 1)[-1] for prefix in prefixes]
+    parts.append(_random_chars())
+    return "-".join(parts)