diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index d3f4821a05b..f0b51c6f1fb 100644 --- a/.github/workflows/test-unit.yml +++ b/.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 diff --git a/tests/unit/proxy/db/test_log_db_metrics_service_spans.py b/tests/unit/proxy/db/test_log_db_metrics_service_spans.py index a7fd830b948..7de9b6fac43 100644 --- a/tests/unit/proxy/db/test_log_db_metrics_service_spans.py +++ b/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) diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 3fb9edc9459..11f4e78619e 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/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: