chore: import upstream snapshot with attribution

sdks/python/tests/conftest.py

@@ -0,0 +1,373 @@
+import os
+import random
+import string
+import tempfile
+import warnings
+from typing import cast
+
+import numpy as np
+
+from unittest import mock
+import pytest
+
+from opik import context_storage
+from opik.api_objects import connection_resources, opik_client
+from opik.healthcheck import connection_monitor, connection_probe
+from opik.message_processing import streamer_constructors
+from opik.message_processing.replay import replay_manager
+
+from . import testlib
+from .testlib import (
+    backend_emulator_message_processor,
+    environment,
+    noop_file_upload_manager,
+)
+
+
+@pytest.fixture(autouse=True)
+def clear_context_storage():
+    yield
+    context_storage.clear_all()
+
+
+@pytest.fixture(autouse=True)
+def shutdown_cached_client_after_test():
+    yield
+    opik_client.reset_global_client(end_client=False)
+    # Drop any shared connection bundle so the process-wide cache doesn't leak a
+    # streamer/emulator from one test into the next. Critical with fake_backend,
+    # which installs a fresh emulator-backed streamer per test: without this, a
+    # reused bundle keeps the previous test's (closed) streamer and the new
+    # emulator records nothing. Fire-and-forget: no flush, no network.
+    connection_resources.MANAGER.close_all(flush=False)
+
+
+@pytest.fixture(autouse=True)
+def skip_local_configuration_file(tmp_path):
+    # Point config path to a non-existent file so the global config is never read
+    # We need to do this because there are a lot of tests that compare the logged project name
+    # in saved traces/spans using "Default Project" as value. But the global config
+    # can already have the project name saved, so the comparison will fail.
+    os.environ["OPIK_CONFIG_PATH"] = str(tmp_path / "nonexistent_opik.config")
+
+
+@pytest.fixture
+def fake_file_upload_manager():
+    upload_manager_emulator = noop_file_upload_manager.FileUploadManagerEmulator()
+    yield upload_manager_emulator
+
+
+@pytest.fixture
+def fake_replay_manager():
+    probe = mock.Mock(spec=connection_probe.ConnectionProbe)
+    probe.check_connection.return_value = connection_probe.ProbeResult(True, None)
+
+    monitor = connection_monitor.OpikConnectionMonitor(
+        ping_interval=10,
+        check_timeout=1,
+        probe=probe,
+    )
+
+    fallback_replay = replay_manager.ReplayManager(
+        monitor=monitor,
+        batch_size=10,
+        batch_replay_delay=0.5,
+        tick_interval_seconds=5.0,
+    )
+    return fallback_replay
+
+
+@pytest.fixture
+def patch_streamer(
+    fake_file_upload_manager: noop_file_upload_manager.FileUploadManagerEmulator,
+    fake_replay_manager: replay_manager.ReplayManager,
+):
+    streamer = None
+    try:
+        # Pass the upload manager to the emulator so it can access attachments
+        fake_message_processor_ = (
+            backend_emulator_message_processor.BackendEmulatorMessageProcessor(
+                file_upload_manager=fake_file_upload_manager
+            )
+        )
+        streamer = streamer_constructors.construct_streamer(
+            message_processor=fake_message_processor_,
+            n_consumers=1,
+            use_batching=True,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            use_attachment_extraction=False,
+            fallback_replay_manager=fake_replay_manager,
+        )
+
+        yield streamer, fake_message_processor_
+    finally:
+        if streamer is not None:
+            # Tests that rely on fake_backend already call `opik.flush_tracker()`
+            # (or access `.trace_trees` after the work is done) during the test
+            # body, so the teardown has nothing meaningful left to drain.
+            streamer.close(flush=False)
+
+
+@pytest.fixture
+def patch_streamer_without_batching(
+    fake_replay_manager: replay_manager.ReplayManager,
+):
+    streamer = None
+    try:
+        # Create an upload manager first
+        fake_upload_manager = noop_file_upload_manager.FileUploadManagerEmulator()
+        # Pass the upload manager to the emulator so it can access attachments
+        fake_message_processor_ = (
+            backend_emulator_message_processor.BackendEmulatorMessageProcessor(
+                file_upload_manager=fake_upload_manager
+            )
+        )
+        streamer = streamer_constructors.construct_streamer(
+            message_processor=fake_message_processor_,
+            n_consumers=1,
+            use_batching=False,
+            file_uploader=fake_upload_manager,
+            max_queue_size=None,
+            use_attachment_extraction=False,
+            fallback_replay_manager=fake_replay_manager,
+        )
+
+        yield streamer, fake_message_processor_
+    finally:
+        if streamer is not None:
+            streamer.close(flush=False)
+
+
+@pytest.fixture
+def fake_backend(patch_streamer):
+    """
+    Patches the function that creates an instance of Streamer under the hood of Opik.
+    As a result, instead of sending data to the backend, it's being passed to
+    the backend emulator, which uses this data to build span and trace trees.
+
+    The resulting trees can be accessed via `fake_backend.trace_trees` or
+    `fake_backend.span_trees` and then used for comparing with expected span/trace trees.
+
+    The trees are built via special classes stored in testlib.models.
+    """
+    streamer, fake_message_processor_ = patch_streamer
+
+    fake_message_processor_ = cast(
+        backend_emulator_message_processor.BackendEmulatorMessageProcessor,
+        fake_message_processor_,
+    )
+
+    mock_construct_online_streamer = mock.Mock()
+    mock_construct_online_streamer.return_value = streamer
+
+    with mock.patch.object(
+        streamer_constructors,
+        "construct_online_streamer",
+        mock_construct_online_streamer,
+    ):
+        yield fake_message_processor_
+
+
+@pytest.fixture
+def fake_backend_without_batching(patch_streamer_without_batching):
+    """
+    Same as fake_backend but must be used when batching is not supported
+    (e.g., when there are Span/Trace update requests)
+    """
+    streamer, fake_message_processor_ = patch_streamer_without_batching
+
+    fake_message_processor_ = cast(
+        backend_emulator_message_processor.BackendEmulatorMessageProcessor,
+        fake_message_processor_,
+    )
+
+    mock_construct_online_streamer = mock.Mock()
+    mock_construct_online_streamer.return_value = streamer
+
+    with mock.patch.object(
+        streamer_constructors,
+        "construct_online_streamer",
+        mock_construct_online_streamer,
+    ):
+        yield fake_message_processor_
+
+
+@pytest.fixture
+def fake_backend_with_patched_environment(request, patch_streamer):
+    """
+    Allows patching environment variables for the duration of the test.
+    Creates a fake backend as in the `fake_backend` fixture.
+    """
+    with testlib.patch_environ(add_keys=request.param):
+        streamer, fake_message_processor_ = patch_streamer
+
+        fake_message_processor_ = cast(
+            backend_emulator_message_processor.BackendEmulatorMessageProcessor,
+            fake_message_processor_,
+        )
+
+        mock_construct_online_streamer = mock.Mock()
+        mock_construct_online_streamer.return_value = streamer
+
+        with mock.patch.object(
+            streamer_constructors,
+            "construct_online_streamer",
+            mock_construct_online_streamer,
+        ):
+            yield fake_message_processor_
+
+
+def random_chars(n: int = 6) -> str:
+    return "".join(random.choice(string.ascii_letters) for _ in range(n))
+
+
+@pytest.fixture()
+def configure_opik_local_env_vars():
+    with testlib.patch_environ(
+        {
+            "OPIK_URL_OVERRIDE": "http://localhost:5173/api",
+        }
+    ):
+        yield
+
+
+@pytest.fixture()
+def configure_opik_not_configured():
+    with testlib.patch_environ(
+        add_keys={},
+        remove_keys=[
+            "OPIK_URL_OVERRIDE",
+            "OPIK_API_KEY",
+            "OPIK_WORKSPACE",
+            "OPIK_PROJECT_NAME",
+        ],
+    ):
+        yield
+
+
+@pytest.fixture()
+def temp_file_15kb():
+    file_size = 15 * 1024
+    with tempfile.NamedTemporaryFile(delete=True) as f:
+        f.write(np.random.bytes(file_size))
+        f.flush()
+
+        yield f
+
+
+@pytest.fixture()
+def temp_file_15mb():
+    file_size = 15 * 1024 * 1024
+    with tempfile.NamedTemporaryFile(delete=True) as f:
+        f.write(np.random.bytes(file_size))
+        f.flush()
+
+        yield f
+
+
+@pytest.fixture(scope="session")
+def ensure_openai_configured():
+    # don't use assertion here to prevent printing os.environ with all env variables
+    if not environment.has_openai_api_key():
+        raise Exception(
+            "OpenAI not configured! Ensure you have set the OPENAI_API_KEY and OPENAI_ORG_ID environment variables."
+        )
+
+
+@pytest.fixture(scope="session")
+def ensure_google_project_and_location_configured():
+    if not (
+        "GOOGLE_CLOUD_PROJECT" in os.environ and "GOOGLE_CLOUD_LOCATION" in os.environ
+    ):
+        raise Exception(
+            "GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION env vars must be set!"
+        )
+
+
+@pytest.fixture(scope="session")
+def ensure_anthropic_configured():
+    # don't use assertion here to prevent printing os.environ with all env variables
+
+    if "ANTHROPIC_API_KEY" not in os.environ:
+        raise Exception("Anthropic not configured!")
+
+
+@pytest.fixture(scope="session")
+def ensure_mistral_configured():
+    # don't use assertion here to prevent printing os.environ with all env variables
+
+    if "MISTRAL_API_KEY" not in os.environ:
+        raise Exception("Mistral not configured!")
+
+
+@pytest.fixture(scope="session")
+def ensure_vertexai_configured(ensure_google_project_and_location_configured):
+    GOOGLE_APPLICATION_CREDENTIALS_PATH = "gcp_credentials.json"
+
+    if "GITHUB_ACTIONS" not in os.environ:
+        if "GOOGLE_APPLICATION_CREDENTIALS" not in os.environ:
+            raise Exception("GOOGLE_APPLICATION_CREDENTIALS env var must be configured")
+        yield
+        return
+
+    if "GCP_CREDENTIALS_JSON" not in os.environ:
+        raise Exception(
+            "GCP_CREDENTIALS_JSON env var with credentials json content must be set"
+        )
+
+    try:
+        gcp_credentials = os.environ["GCP_CREDENTIALS_JSON"]
+        with open(GOOGLE_APPLICATION_CREDENTIALS_PATH, mode="wt") as output_file:
+            output_file.write(gcp_credentials)
+
+        with testlib.patch_environ(
+            add_keys={
+                "GOOGLE_APPLICATION_CREDENTIALS": GOOGLE_APPLICATION_CREDENTIALS_PATH
+            }
+        ):
+            yield
+    finally:
+        if os.path.exists(GOOGLE_APPLICATION_CREDENTIALS_PATH):
+            os.remove(GOOGLE_APPLICATION_CREDENTIALS_PATH)
+
+
+@pytest.fixture(scope="session")
+def ensure_google_api_configured():
+    GOOGLE_API_KEY = "GOOGLE_API_KEY"
+    if GOOGLE_API_KEY not in os.environ:
+        raise Exception(f"{GOOGLE_API_KEY} env var must be set")
+
+
+@pytest.fixture(scope="session")
+def ensure_aws_bedrock_configured():
+    import boto3
+
+    session = boto3.Session()
+
+    bedrock_client = session.client(service_name="bedrock")
+    try:
+        available_models = bedrock_client.list_foundation_models()["modelSummaries"]
+        if not available_models:
+            raise Exception("AWS Bedrock not configured! No models available")
+    except Exception as e:
+        raise Exception(f"AWS Bedrock not configured! {e}")
+
+
+@pytest.fixture(scope="session")
+def ensure_groq_configured():
+    # don't use assertion here to prevent printing os.environ with all env variables
+    if "GROQ_API_KEY" not in os.environ:
+        raise Exception("Groq is not configured!")
+
+
+warnings.filterwarnings(
+    "ignore",
+    message=".*PydanticDeprecatedSince20.*",
+    category=DeprecationWarning,
+)
+warnings.filterwarnings(
+    "ignore",
+    message=".*0 counts of 2-gram overlaps.*",
+    category=UserWarning,
+)