From 83e2e5bf99bf443fa6b4723be6a49c39503b880d Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 7 Nov 2026 23:37:18 +0000
Subject: [PATCH] test(integration): budget proxy restart tests from calibrated
 boot time (#49761)

Both restart tests boot two 2-worker proxies but capped the case at a fixed
180s, below even one calibrated 2-worker readiness budget. On a loaded runner
each boot took 80-90s, so pytest-timeout fired mid-boot or mid-shutdown

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../test_azure_content_safety_dispatch_resilience.py         | 4 +++-
 tests/integration/observability/test_azure_storage_chaos.py  | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py
index 14eb687f69f..d184112f1ec 100644
--- a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py
+++ b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py
@@ -8,11 +8,13 @@ from typing import Final
 
 import httpx
 import pytest
+from integration._support.budget import readiness_seconds
 from integration._support.client import Gateway, Scenario, eventually
 from integration._support.database import read_rows
 from integration._support.redis_process import owned_redis
 from integration._support.wire import Wire, wire_server
 from integration._support.ports import allocate_port
+from integration._support.process import graceful_stop_seconds
 from integration.observability.azure_dispatch_support import (
     ATTACK_MARKER,
     AzureBehavior,
@@ -406,7 +408,7 @@ def test_x3_killing_one_worker_leaves_the_other_serving(
                 _assert_spend(response.headers["x-litellm-call-id"], response.text)
 
 
-@pytest.mark.timeout(180)
+@pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 120)
 def test_x4_proxy_restart_preserves_completed_spend_rows(
     gateway: Gateway, tmp_path: Path
 ) -> None:
diff --git a/tests/integration/observability/test_azure_storage_chaos.py b/tests/integration/observability/test_azure_storage_chaos.py
index acc2eb1439d..e78f1a4a1db 100644
--- a/tests/integration/observability/test_azure_storage_chaos.py
+++ b/tests/integration/observability/test_azure_storage_chaos.py
@@ -13,8 +13,9 @@ from _azure_storage_support import (
     collect_files,
 )
 from _s3_v2_support import matched_ids, mixed_burst, surface_reply
+from integration._support.budget import readiness_seconds
 from integration._support.client import Gateway, JsonValue, eventually
-from integration._support.process import group_members, owned_proxy_process
+from integration._support.process import graceful_stop_seconds, group_members, owned_proxy_process
 from integration._support.tls import server_context, write_self_signed_cert
 from integration._support.wire import wire_server
 
@@ -175,7 +176,7 @@ def test_killing_one_worker_keeps_the_other_serving_and_uploading(gateway: Gatew
             )
 
 
-@pytest.mark.timeout(180)
+@pytest.mark.timeout(2 * readiness_seconds(WORKERS) + 2 * graceful_stop_seconds() + 120)
 def test_restarting_the_proxy_before_the_queue_flushes_bounds_the_loss_to_the_unflushed_queue_and_recovers(
     gateway: Gateway, tmp_path: Path
 ) -> None: