test: inject clocks and wait on conditions in three wall-clock-flaky unit tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/unit/integrations/generic_api/test_generic_api_callback.py

@@ -29,6 +29,13 @@ def set_verbose_logger_level() -> Iterator[None]:
     verbose_logger.setLevel(original_level)
 
 
+async def _wait_for_events_and_flush(generic_logger: GenericAPILogger, expected: int) -> None:
+    async with asyncio.timeout(10):
+        while len(generic_logger.log_queue) < expected:
+            await asyncio.sleep(0.01)
+    await generic_logger.flush_queue()
+
+
 @pytest.mark.asyncio
 async def test_generic_api_callback():
     """
@@ -48,7 +55,7 @@ async def test_generic_api_callback():
     os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint
 
     # Initialize the GenericAPILogger and set the mock
-    generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=1)
+    generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=3600)
     generic_logger.async_httpx_client.post = mock_post
     litellm.callbacks = [generic_logger]
 
@@ -60,8 +67,8 @@ async def test_generic_api_callback():
         user="test_user",
     )
 
-    # Wait for async flush
-    await asyncio.sleep(3)
+    # Wait for this test's event to reach the queue, then flush it
+    await _wait_for_events_and_flush(generic_logger, 1)
 
     # Assert httpx post was called
     mock_post.assert_called_once()
@@ -127,7 +134,7 @@ async def test_generic_api_callback_multiple_logs():
     os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint
 
     # Initialize the GenericAPILogger and set the mock
-    generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=5)
+    generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=3600)
     generic_logger.async_httpx_client.post = mock_post
     litellm.callbacks = [generic_logger]
 
@@ -141,7 +148,7 @@ async def test_generic_api_callback_multiple_logs():
         )
 
     # Wait for async flush
-    await asyncio.sleep(6)
+    await _wait_for_events_and_flush(generic_logger, 10)
 
     # Assert httpx post was called
     mock_post.assert_called_once()
@@ -205,7 +212,7 @@ async def test_generic_api_callback_ndjson_format():
     generic_logger = GenericAPILogger(
         endpoint=test_endpoint,
         headers=test_headers,
-        flush_interval=1,
+        flush_interval=3600,
         log_format="ndjson",  # Set NDJSON format
     )
     generic_logger.async_httpx_client.post = mock_post
@@ -221,7 +228,7 @@ async def test_generic_api_callback_ndjson_format():
         )
 
     # Wait for async flush
-    await asyncio.sleep(3)
+    await _wait_for_events_and_flush(generic_logger, 3)
 
     # Assert httpx post was called
     mock_post.assert_called_once()
@@ -274,7 +281,7 @@ async def test_generic_api_callback_single_format():
     generic_logger = GenericAPILogger(
         endpoint=test_endpoint,
         headers=test_headers,
-        flush_interval=1,  # Quick flush to trigger batch send
+        flush_interval=3600,
         log_format="single",  # Set single format
     )
     generic_logger.async_httpx_client.post = mock_post
@@ -290,7 +297,7 @@ async def test_generic_api_callback_single_format():
         )
 
     # Wait for async flush
-    await asyncio.sleep(3)
+    await _wait_for_events_and_flush(generic_logger, 3)
 
     # Assert httpx post was called 3 times (once per log in batch)
     assert mock_post.call_count == 3, f"Expected 3 calls, got {mock_post.call_count}"
@@ -333,7 +340,7 @@ async def test_generic_api_callback_json_array_format_explicit():
     generic_logger = GenericAPILogger(
         endpoint=test_endpoint,
         headers=test_headers,
-        flush_interval=1,
+        flush_interval=3600,
         log_format="json_array",  # Explicitly set json_array
     )
     generic_logger.async_httpx_client.post = mock_post
@@ -349,7 +356,7 @@ async def test_generic_api_callback_json_array_format_explicit():
         )
 
     # Wait for async flush
-    await asyncio.sleep(3)
+    await _wait_for_events_and_flush(generic_logger, 5)
 
     # Assert httpx post was called once (batched)
     mock_post.assert_called_once()
@@ -383,7 +390,7 @@ async def test_generic_api_callback_sumologic_uses_ndjson():
     os.environ["SUMOLOGIC_WEBHOOK_URL"] = "https://collectors.sumologic.com/receiver/v1/http/test123"
 
     # Initialize using callback_name (loads from JSON config)
-    generic_logger = GenericAPILogger(callback_name="sumologic", flush_interval=1)
+    generic_logger = GenericAPILogger(callback_name="sumologic", flush_interval=3600)
     generic_logger.async_httpx_client.post = mock_post
     litellm.callbacks = [generic_logger]
 
@@ -400,7 +407,7 @@ async def test_generic_api_callback_sumologic_uses_ndjson():
         )
 
     # Wait for async flush
-    await asyncio.sleep(3)
+    await _wait_for_events_and_flush(generic_logger, 2)
 
     # Assert httpx post was called
     mock_post.assert_called_once()

tests/unit/integrations/langfuse/test_langfuse_sdk.py

@@ -1080,11 +1080,19 @@ def test_auth_check_and_project_id_make_one_round_trip_when_langfuse_is_down(
     [(500, 2), (503, 2), (429, 1), (404, 1)],
     ids=["http-500", "http-503", "http-429", "http-404"],
 )
-def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_trips: int):
+def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(
+    status: int, round_trips: int, monkeypatch: pytest.MonkeyPatch
+) -> None:
     """A cold ``get_prompt`` fetches inline on the event loop; with the generated client's default retries a
     429 carrying ``Retry-After: 30`` used to hold the loop for a minute. A 5xx gets the v2 client's one
     quick retry, a 429 or 4xx none."""
     requests: list[httpx.Request] = []
+    sleeps: Final[list[float]] = []
+    monkeypatch.setattr(
+        langfuse_http_client,
+        "time",
+        SimpleNamespace(sleep=sleeps.append, time=time.time),
+    )
 
     def fail(request: httpx.Request) -> httpx.Response:
         requests.append(request)
@@ -1097,18 +1105,25 @@ def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_
         httpx_client=httpx.Client(transport=httpx.MockTransport(fail)),
     )
 
-    started = monotonic()
     with pytest.raises(LangfusePromptError) as caught:
         client.get_prompt("greeting")
     assert len(requests) == round_trips
-    assert monotonic() - started < 0.5
+    assert sleeps == [], "a cold prompt miss should not sleep for REST retries"
     assert caught.value.status_code == status
 
 
 @pytest.mark.parametrize("first_failure", [503, "connect-error"], ids=["http-503", "connect-error"])
-def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(first_failure: int | str):
+def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(
+    first_failure: int | str, monkeypatch: pytest.MonkeyPatch
+) -> None:
     """The v2 client retried a cold fetch once; a single Langfuse blip must not fail the LLM call."""
     requests: list[httpx.Request] = []
+    sleeps: Final[list[float]] = []
+    monkeypatch.setattr(
+        langfuse_http_client,
+        "time",
+        SimpleNamespace(sleep=sleeps.append, time=time.time),
+    )
 
     def flaky(request: httpx.Request) -> httpx.Response:
         requests.append(request)
@@ -1125,10 +1140,9 @@ def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(firs
         httpx_client=httpx.Client(transport=httpx.MockTransport(flaky)),
     )
 
-    started = monotonic()
     assert client.get_prompt("greeting").compile() == "hello"
     assert len(requests) == 2
-    assert monotonic() - started < 0.5
+    assert sleeps == [], "one transient failure should not sleep for REST retries"
     assert client.get_prompt("greeting").compile() == "hello", "the retried prompt is cached like any other"
     assert len(requests) == 2
 

tests/unit/proxy/hooks/test_autorouter_baseline_cache.py

@@ -499,12 +499,14 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
     from litellm.litellm_core_utils import logging_worker
     from litellm.litellm_core_utils.logging_worker import LoggingWorker
 
-    release: Final = asyncio.Event()
+    counted: Final = asyncio.Event()
+    loop: Final = asyncio.get_running_loop()
+    real_time: Final = loop.time
+    monkeypatch.setattr(loop, "time", lambda: real_time() + (2.05 if counted.is_set() else 0.0))
 
     async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int:
-        if not release.is_set():
-            asyncio.get_running_loop().call_later(2.05, release.set)
-            await release.wait()
+        counted.set()
+        await asyncio.sleep(0.001)
         return await _count(model, api_key, body)
 
     worker: Final = LoggingWorker(timeout=8.0)
@@ -519,7 +521,6 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
         assert observed.plan is not None and observed.plan.breakpoints[0].prefix_tokens == 5000
         assert payload["response_cost"] is not None and worker._timeout_total == 0
     finally:
-        release.set()
         await worker.stop()