chore: import upstream snapshot with attribution

sdks/python/tests/unit/api_objects/test_search_helpers.py

@@ -0,0 +1,216 @@
+import json
+import logging
+from unittest import mock
+
+import pytest
+
+from opik.api_objects import search_helpers
+from opik.rest_api.core.api_error import ApiError
+from opik.rest_api.types import span_public, trace_public
+
+
+SPAN_RECORD = {
+    "id": "0195f6f1-9da2-7630-b285-60cf5580372f",
+    "project_id": "0195f6f1-9c82-751f-a1ad-54fec4b5c7d8",
+    "trace_id": "0195f6f1-9c3f-7b65-a334-56173d19bc00",
+    "name": "synthesize",
+    "type": "general",
+    "start_time": "2025-01-03T11:03:17.875608Z",
+    "end_time": "2025-01-03T11:03:18.591814Z",
+    "input": {"q": "x"},
+    "output": {"a": "y"},
+    "created_at": "2025-04-02T14:39:44.412550Z",
+    "last_updated_at": "2025-04-02T14:39:44.412550Z",
+    "created_by": "admin",
+    "last_updated_by": "admin",
+    "duration": 1.0,
+}
+
+TRACE_RECORD = {
+    "id": "0195f6f1-9c3f-7b65-a334-56173d19bc00",
+    "project_id": "0195f6f1-9c82-751f-a1ad-54fec4b5c7d8",
+    "name": "root",
+    "start_time": "2025-01-03T11:03:17.505783Z",
+    "end_time": "2025-01-03T11:03:18.591897Z",
+    "input": {"q": "x"},
+    "output": {"a": "y"},
+    "created_at": "2025-04-02T14:39:44.412550Z",
+    "last_updated_at": "2025-04-02T14:39:44.412550Z",
+    "created_by": "admin",
+    "last_updated_by": "admin",
+    "duration": 1.0,
+}
+
+
+def _stream(record: dict):
+    """Mimic the bytes-iterator returned by the Fern-generated search clients."""
+    yield f"{json.dumps(record)}\r\n".encode("utf-8")
+
+
+@pytest.fixture
+def patched_sleep():
+    # Patch the rest_helpers retry-delay indirection, NOT the global time.sleep:
+    # a global patch turns background daemon threads' pacing sleep into a busy
+    # loop that can starve the interpreter and hang the suite.
+    with mock.patch("opik.api_objects.rest_helpers._sleep") as sleep_mock:
+        yield sleep_mock
+
+
+def _sleep_called_with(mock_sleep: mock.Mock, seconds) -> bool:
+    return any(call.args == (seconds,) for call in mock_sleep.call_args_list)
+
+
+def test_search_spans__429_response__retries_with_reset_header(
+    patched_sleep, capture_log
+):
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.side_effect = [
+        ApiError(status_code=429, headers={"RateLimit-Reset": "2"}),
+        _stream(SPAN_RECORD),
+    ]
+
+    result = search_helpers.search_spans_with_filters(
+        rest_client=rest_client,
+        trace_id=None,
+        project_name="proj",
+        filters=None,
+        max_results=10,
+        truncate=False,
+    )
+
+    assert len(result) == 1
+    assert result[0] == span_public.SpanPublic.model_validate(SPAN_RECORD)
+    assert rest_client.spans.search_spans.call_count == 2
+    assert _sleep_called_with(patched_sleep, 2)
+    assert any(
+        "search_spans" in record.message and record.levelno == logging.WARNING
+        for record in capture_log.records
+    )
+
+
+def test_search_traces__429_response__retries_with_reset_header(
+    patched_sleep, capture_log
+):
+    rest_client = mock.Mock()
+    rest_client.traces.search_traces.side_effect = [
+        ApiError(status_code=429, headers={"RateLimit-Reset": "3"}),
+        _stream(TRACE_RECORD),
+    ]
+
+    result = search_helpers.search_traces_with_filters(
+        rest_client=rest_client,
+        project_name="proj",
+        filters=None,
+        max_results=10,
+        truncate=False,
+    )
+
+    assert len(result) == 1
+    assert result[0] == trace_public.TracePublic.model_validate(TRACE_RECORD)
+    assert rest_client.traces.search_traces.call_count == 2
+    assert _sleep_called_with(patched_sleep, 3)
+    assert any(
+        "search_traces" in record.message and record.levelno == logging.WARNING
+        for record in capture_log.records
+    )
+
+
+def test_search_spans__non_429_error__propagates_without_retry():
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.side_effect = ApiError(
+        status_code=400, body="bad request"
+    )
+
+    with pytest.raises(ApiError) as exc_info:
+        search_helpers.search_spans_with_filters(
+            rest_client=rest_client,
+            trace_id=None,
+            project_name="proj",
+            filters=None,
+            max_results=10,
+            truncate=False,
+        )
+
+    assert exc_info.value.status_code == 400
+    # call_count == 1 already proves no retry happened — no need to assert on
+    # patched_sleep, which is polluted by background-thread sleeps.
+    assert rest_client.spans.search_spans.call_count == 1
+
+
+def test_search_traces__non_429_error__propagates_without_retry():
+    rest_client = mock.Mock()
+    rest_client.traces.search_traces.side_effect = ApiError(
+        status_code=500, body="boom"
+    )
+
+    with pytest.raises(ApiError) as exc_info:
+        search_helpers.search_traces_with_filters(
+            rest_client=rest_client,
+            project_name="proj",
+            filters=None,
+            max_results=10,
+            truncate=False,
+        )
+
+    assert exc_info.value.status_code == 500
+    assert rest_client.traces.search_traces.call_count == 1
+
+
+def test_search_spans__max_batch_size__passed_as_limit_to_rest_client():
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.return_value = _stream(SPAN_RECORD)
+
+    search_helpers.search_spans_with_filters(
+        rest_client=rest_client,
+        trace_id=None,
+        project_name="proj",
+        filters=None,
+        max_results=10_000,
+        truncate=False,
+        max_batch_size=400,
+    )
+
+    # max_results (10_000) is the total cap; max_batch_size (400) is the
+    # per-request page size that reaches the Fern client as ``limit``.
+    _, kwargs = rest_client.spans.search_spans.call_args
+    assert kwargs["limit"] == 400
+
+
+def test_search_traces__max_batch_size__passed_as_limit_to_rest_client():
+    rest_client = mock.Mock()
+    rest_client.traces.search_traces.return_value = _stream(TRACE_RECORD)
+
+    search_helpers.search_traces_with_filters(
+        rest_client=rest_client,
+        project_name="proj",
+        filters=None,
+        max_results=10_000,
+        truncate=False,
+        max_batch_size=250,
+    )
+
+    _, kwargs = rest_client.traces.search_traces.call_args
+    assert kwargs["limit"] == 250
+
+
+def test_search_spans__429_without_reset_header__falls_back_to_one_second_sleep(
+    patched_sleep,
+):
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.side_effect = [
+        ApiError(status_code=429, headers={}),
+        _stream(SPAN_RECORD),
+    ]
+
+    result = search_helpers.search_spans_with_filters(
+        rest_client=rest_client,
+        trace_id=None,
+        project_name="proj",
+        filters=None,
+        max_results=10,
+        truncate=False,
+    )
+
+    assert len(result) == 1
+    assert rest_client.spans.search_spans.call_count == 2
+    assert _sleep_called_with(patched_sleep, 1)