diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 5610ac8fff6..318ddd51c38 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -279,14 +279,16 @@ def mock_get(key, default=None): "response_cost": 0.0, } + # Use fixed timestamps to avoid timing-related flakiness + fixed_time = datetime.datetime(2024, 1, 1, 12, 0, 0) # Call the method under test self.logger._log_langfuse_v2( user_id="test-user", metadata={}, litellm_params=kwargs["litellm_params"], output={"role": "assistant", "content": "Response"}, - start_time=datetime.datetime.now(), - end_time=datetime.datetime.now(), + start_time=fixed_time, + end_time=fixed_time + datetime.timedelta(seconds=1), kwargs=kwargs, optional_params=kwargs["optional_params"], input={"messages": kwargs["messages"]}, diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py index fc45a13c2c1..29cafd4f930 100644 --- a/tests/test_litellm/llms/watsonx/test_watsonx.py +++ b/tests/test_litellm/llms/watsonx/test_watsonx.py @@ -279,7 +279,11 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): ), f"POST should have been called at least once, got {mock_post.call_count}" # Get the request body from the first call - call_args = mock_post.call_args + # Use call_args_list to be more robust - get the first call's arguments + assert len(mock_post.call_args_list) > 0, "mock_post should have at least one call" + call_args = mock_post.call_args_list[0] + assert call_args is not None, "call_args should not be None" + assert "data" in call_args.kwargs, "call_args.kwargs should contain 'data'" json_data = json.loads(call_args.kwargs["data"]) print(f"\n{'='*80}") diff --git a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py b/tests/test_litellm/responses/test_no_duplicate_spend_logs.py index 745dc782239..249374b3078 100644 --- a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py +++ b/tests/test_litellm/responses/test_no_duplicate_spend_logs.py @@ -83,9 +83,10 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti mock_response="Hello! I'm doing well." # Use mock to avoid real API call ) - # Give async logging time to complete - import asyncio - await asyncio.sleep(1) + # Wait for async logging to complete using the logging worker's flush method + # This is more reliable than sleep() which can cause race conditions + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + await GLOBAL_LOGGING_WORKER.flush() # Verify that log_success_event was called exactly once assert spend_logger.log_count == 1, (