test: deflake proxy metrics and shutdown drain (#51078)

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

.github/workflows/test-unit.yml

@@ -633,23 +633,32 @@ jobs:
             job-timeout-minutes: 60
             dist: loadscope
 
-          - shard: proxy-infra
+          - shard: proxy-infra-db-spend
             test-path: |-
               tests/unit/proxy/db
               --ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py
               --ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py
-              tests/unit/proxy/middleware
-              --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py
               tests/unit/proxy/spend_tracking
               --ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py
-              tests/unit/proxy/pass_through_endpoints
-              tests/unit/proxy/mcp_server/test_tool_registry.py
               tests/unit/proxy/common_utils
               --ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py
               --ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py
               --ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py
               --ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py
               --ignore=tests/unit/proxy/common_utils/test_realtime_cache.py
+            python-version: "3.12"
+            workers: 4
+            reruns: 0
+            timeout-minutes: 20
+            job-timeout-minutes: 60
+            dist: loadscope
+
+          - shard: proxy-infra-routing
+            test-path: |-
+              tests/unit/proxy/middleware
+              --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py
+              tests/unit/proxy/pass_through_endpoints
+              tests/unit/proxy/mcp_server/test_tool_registry.py
               tests/unit/proxy/enterprise_billing
               tests/unit/proxy/types_utils
               tests/unit/proxy/logging_endpoints

tests/unit/proxy/db/test_log_db_metrics_service_spans.py

@@ -18,12 +18,12 @@ from prisma.errors import ClientNotConnectedError
 from pydantic import TypeAdapter
 
 import litellm
-from litellm._service_logger import ServiceTypes
+from litellm._service_logger import ServiceLogging, ServiceTypes
 from litellm.integrations.datadog.datadog import DataDogLogger
 from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
 from litellm.proxy.db.log_db_metrics import log_db_metrics
 from litellm.proxy.db.prisma_client import _PrismaDrainTracker, _TrackedPrismaEngine
-from litellm.proxy.proxy_server import proxy_logging_obj
+from litellm.proxy import proxy_server
 from litellm.integrations.prometheus_services import PrometheusServicesLogger
 from prometheus_client import REGISTRY
 
@@ -80,10 +80,10 @@ def rig(monkeypatch: pytest.MonkeyPatch) -> _Rig:
     with patch("asyncio.create_task", side_effect=_discard_periodic_flush):
         datadog: Final = DataDogLogger()
     monkeypatch.setattr(litellm, "service_callback", [otel, datadog, "prometheus_system"])
-    monkeypatch.setattr(proxy_logging_obj.service_logging_obj, "dd_logger", datadog, raising=False)
-    monkeypatch.setattr(
-        proxy_logging_obj.service_logging_obj, "prometheusServicesLogger", PrometheusServicesLogger(), raising=False
-    )
+    service_logging: Final = ServiceLogging()
+    service_logging.dd_logger = datadog
+    service_logging.prometheusServicesLogger = PrometheusServicesLogger()
+    monkeypatch.setattr(proxy_server.proxy_logging_obj, "service_logging_obj", service_logging)
     return _Rig(exporter, provider, datadog)
 
 

tests/unit/responses/test_streaming_iterator.py

@@ -2743,14 +2743,17 @@ async def test_shutdown_drain_waits_for_a_pending_logging_task():
 
 @pytest.mark.asyncio
 async def test_shutdown_drain_gives_up_on_a_stuck_logging_task_at_its_timeout():
-    stuck: Final = streaming_module.spawn_logging_task(
-        asyncio.get_running_loop(), asyncio.Event().wait(), task_name="stuck"
-    )
+    loop: Final = asyncio.get_running_loop()
+    with patch.object(loop, "time", return_value=0.0):
+        stuck: Final = streaming_module.spawn_logging_task(loop, asyncio.Event().wait(), task_name="stuck")
 
-    await asyncio.wait_for(streaming_module.drain_logging_tasks(timeout=0.05), timeout=1)
-
-    assert not stuck.done()
-    stuck.cancel()
+        try:
+            await streaming_module.drain_logging_tasks(timeout=0)
+            assert not stuck.done()
+        finally:
+            stuck.cancel()
+            with suppress(asyncio.CancelledError):
+                await stuck
 
 
 async def _drain_held_until_released(release: asyncio.Event) -> bool: