diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_modes.py b/tests/integration/observability/test_arize_otel_v2_openinference_modes.py index cb61ccf2ac1..51709cbd713 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_modes.py +++ b/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")) diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index 48e9930d948..be38d5a6f1f 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/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 diff --git a/tests/integration/observability/test_otel_team_capture_message_content.py b/tests/integration/observability/test_otel_team_capture_message_content.py index 84da9a40486..f8f2b85bab5 100644 --- a/tests/integration/observability/test_otel_team_capture_message_content.py +++ b/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) diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py index 382463eb46c..a81f21f9839 100644 --- a/tests/integration/observability/test_s3_v2_partition_granularity.py +++ b/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: diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 33904637366..8c4ea0d4136 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/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,