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)