diff --git a/tests/local_testing/test_opik.py b/tests/local_testing/test_opik.py index 8be4b796360..2f6b15e1f27 100644 --- a/tests/local_testing/test_opik.py +++ b/tests/local_testing/test_opik.py @@ -16,6 +16,8 @@ verbose_logger.setLevel(logging.DEBUG) litellm.set_verbose = True import time +INTERVAL_TOO_LONG_TO_FIRE_DURING_THIS_TEST = 3600 + @pytest.mark.asyncio async def test_opik_logging_http_request(): @@ -23,70 +25,60 @@ async def test_opik_logging_http_request(): - Test that HTTP requests are made to Opik - Traces and spans are batched correctly """ - try: - from litellm.integrations.opik.opik import OpikLogger + from litellm.integrations.opik.opik import OpikLogger - os.environ["OPIK_URL_OVERRIDE"] = "https://fake.comet.com/opik/api" - os.environ["OPIK_API_KEY"] = "anything" - os.environ["OPIK_WORKSPACE"] = "anything" + os.environ["OPIK_URL_OVERRIDE"] = "https://fake.comet.com/opik/api" + os.environ["OPIK_API_KEY"] = "anything" + os.environ["OPIK_WORKSPACE"] = "anything" - # Initialize OpikLogger - test_opik_logger = OpikLogger() + test_opik_logger = OpikLogger() + test_opik_logger.flush_interval = INTERVAL_TOO_LONG_TO_FIRE_DURING_THIS_TEST + test_opik_logger.batch_size = 12 - litellm.callbacks = [test_opik_logger] - test_opik_logger.batch_size = 12 - litellm.set_verbose = True + litellm.callbacks = [test_opik_logger] - # Create a mock for the async_client's post method - mock_post = AsyncMock() - mock_post.return_value.status_code = 202 - mock_post.return_value.text = "Accepted" - test_opik_logger.async_httpx_client.post = mock_post + mock_post = AsyncMock(return_value=Mock(status_code=202, text="Accepted")) + test_opik_logger.async_httpx_client.post = mock_post - # Make multiple calls to ensure we don't hit the batch size - for _ in range(5): - response = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Test message"}], - max_tokens=10, - temperature=0.2, - mock_response="This is a mock response", - ) - await asyncio.sleep(1) + def opik_batch_calls(): + return [ + call + for call in mock_post.call_args_list + if "/traces/batch" in str(call) or "/spans/batch" in str(call) + ] - # Check batching of events and that the queue contains 5 trace events and 5 span events - assert ( - mock_post.called == False - ), "HTTP request was made but events should have been batched" - assert len(test_opik_logger.log_queue) == 10 + for _ in range(5): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + ) + await asyncio.sleep(1) - # Now make calls to exceed the batch size - for _ in range(3): - response = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Test message"}], - max_tokens=10, - temperature=0.2, - mock_response="This is a mock response", - ) + assert opik_batch_calls() == [], "events below batch_size must stay queued" + assert len(test_opik_logger.log_queue) == 10 - # Wait a short time for any asynchronous operations to complete - await asyncio.sleep(1) + for _ in range(3): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + ) + await asyncio.sleep(1) - # Check that the queue was flushed after exceeding batch size - assert len(test_opik_logger.log_queue) < test_opik_logger.batch_size + assert opik_batch_calls(), "crossing batch_size must flush the queue" + events_left_over_after_the_size_triggered_flush = len(test_opik_logger.log_queue) + assert 0 < events_left_over_after_the_size_triggered_flush < test_opik_logger.batch_size - # Check that the data has been sent when it goes above the flush interval - await asyncio.sleep(test_opik_logger.flush_interval) - assert len(test_opik_logger.log_queue) == 0 + calls_before_periodic_flush = len(opik_batch_calls()) + await test_opik_logger.flush_queue() - # Clean up - for cb in litellm.callbacks: - if isinstance(cb, OpikLogger): - await cb.async_httpx_client.client.aclose() - - except Exception as e: - pytest.fail(f"Error occurred: {e}") + assert len(opik_batch_calls()) > calls_before_periodic_flush + assert len(test_opik_logger.log_queue) == 0 def test_sync_opik_logging_http_request():