Merge pull request #49486 from BerriAI/litellm_chore_935b3f

test: backport three test fixes from main to rc/1.108.0

tests/integration/observability/test_arize_otel_v2_openinference_modes.py

@@ -24,6 +24,7 @@ from _openinference_support import (
     _rig,
 )
 from integration._support.client import Gateway
+from integration._support.process import graceful_stop_seconds
 from integration._support.wire import Reply, Request
 
 _DEFAULT_METADATA: Final = {
@@ -411,7 +412,7 @@ def test_arize_otel_v2_c9_upstream_failures_are_recorded(failure_status: int, ga
         assert not any(".tool_calls." in key for key in attributes), attributes
 
 
-@pytest.mark.timeout(180)
+@pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
 def test_arize_otel_v2_c10_attribute_limit_keeps_tool_prefix(gateway: Gateway, tmp_path: Path) -> None:
     marker: Final = "c10-" + uuid.uuid4().hex
     calls: Final = _response_tool_calls(marker, ("Paris", "Berlin", "Rome", "Tokyo", "Oslo", "Lima", "Accra", "Delhi"))

tests/integration/observability/test_logging_only_scope_chaos.py

@@ -42,7 +42,7 @@ from _logging_only_scope_support import (
 )
 from anthropic import APIConnectionError as AnthropicAPIConnectionError
 from integration._support.client import Gateway, eventually
-from integration._support.process import owned_proxy, owned_proxy_process
+from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process
 from integration._support.timing import Timed, assert_not_delayed_by, timed
 from integration._support.wire import Reply, Request
 from integration._support.ports import allocate_port
@@ -52,7 +52,7 @@ from pydantic import JsonValue
 _EDGE_DELAY_SECONDS: Final = 5
 
 
-@pytest.mark.timeout(180)
+@pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
 def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gateway: Gateway, tmp_path: Path) -> None:
     identity: Final = f"logging-scope-k1-{uuid.uuid4().hex}"
     marker: Final = uuid.uuid4().hex
@@ -210,7 +210,7 @@ def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gat
                 ) == tuple((identity, "logging_only", expected_status) for _ in expected_directions), (index, entries)
 
 
-@pytest.mark.timeout(180)
+@pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
 def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway, tmp_path: Path) -> None:
     identity: Final = f"logging-scope-k2-{uuid.uuid4().hex}"
     marker: Final = uuid.uuid4().hex

tests/integration/observability/test_otel_team_capture_message_content.py

@@ -302,6 +302,23 @@ def _with_content(url: str, since: int, sent: Sent) -> tuple[Span, ...]:
     )
 
 
+def _is_root(span: Span) -> bool:
+    return span["kind"] == SERVER and span["name"].startswith("POST")
+
+
+def _with_root(url: str, since: int, sent: Sent) -> tuple[Span, ...]:
+    """The copy of ``sent`` at ``url`` once its model span, the marker and its server span arrived."""
+    return eventually(
+        lambda: _of_request(url, since, sent),
+        lambda spans: (
+            bool(_model_spans(spans))
+            and any(_carries(span, sent.marker) for span in spans)
+            and any(_is_root(span) for span in spans)
+        ),
+        seconds=40,
+    )
+
+
 def _arrived(
     url: str, since: int, sent: Sent, expected: frozenset[str], credential: Credential | None = None
 ) -> tuple[Span, ...]:
@@ -605,11 +622,7 @@ def test_stalled_tenant_during_a_mixed_burst_keeps_serving_and_each_teams_policy
         configure_sink(rig.sinks.tenant, paused=False)
     assert [result.status for result in sent] == [200] * 30, [result.text for result in sent if result.status != 200]
     for (key, _, _), result in zip(plan, sent, strict=True):
-        roots = [
-            span
-            for span in _with_content(rig.sinks.operator, cursors.operator, result)
-            if span["kind"] == SERVER and span["name"].startswith("POST ")
-        ]
+        roots = [span for span in _with_root(rig.sinks.operator, cursors.operator, result) if _is_root(span)]
         assert len(roots) == 1, roots
         if key == restricted:
             _assert_redacted_twin(rig, result, cursors, rig.sinks.tenant, cursors.tenant)

tests/integration/observability/test_s3_v2_partition_granularity.py

@@ -30,7 +30,7 @@ from _s3_v2_support import (
 from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value
 from integration._support.database import read_rows, scratch_database
 from integration._support.database_relay import database_relay
-from integration._support.process import OwnedProxy, owned_proxy_process
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
 from integration._support.wire import Reply, Request, wire_server
 
 FLUSH: Final = {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"}
@@ -1158,7 +1158,7 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway
     assert _outside_layout(objects, "hour") == ()
 
 
-@pytest.mark.timeout(180)
+@pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
 def test_s3_v2_hour_proxy_restart_mid_burst_keeps_the_layout_without_duplicates(
     gateway: Gateway, tmp_path: Path
 ) -> None:

tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py

@@ -10,7 +10,6 @@ import pytest
 import litellm
 from litellm.caching.caching import DualCache
 from litellm.caching.redis_cache import RedisCache
-from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
 from litellm.proxy.hooks.model_max_budget_limiter import (
     PROXY_VirtualKeyModelMaxBudgetLimiter,
@@ -20,6 +19,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
     model_budget_start_time_cache_key,
     resolve_model_budget,
 )
+from litellm.router_strategy.budget_limiter import budget_window_seconds
 from litellm.types.caching import RedisPipelineIncrementOperation
 
 
@@ -31,7 +31,7 @@ async def _reset_at_of_open_window(
             entity_type=entity_type, entity_id=entity_id, budget_model=budget_model, budget_duration=period
         )
     )
-    return datetime.fromtimestamp(started_at + duration_in_seconds(period), tz=timezone.utc).isoformat()
+    return datetime.fromtimestamp(started_at + budget_window_seconds(started_at, period), tz=timezone.utc).isoformat()
 
 
 # Test class setup
@@ -973,7 +973,7 @@ async def test_user_model_budget_window_resets_when_the_period_elapses():
 
     # Age the window past its period. The next charge opens a new window rather
     # than adding to the exhausted one.
-    elapsed = duration_in_seconds("1mo") + 60
+    elapsed = int(timedelta(days=31, minutes=1).total_seconds())
     await dual_cache.async_set_cache(
         key=start_time_key,
         value=datetime.now(timezone.utc).timestamp() - elapsed,