From 369d3244fbf13b9286cf00b7ea349f324f001ee1 Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Sat, 31 Oct 2026 10:17:34 -0700
Subject: [PATCH] test(integration): size restart-mid-burst timeouts from the
 proxy boot budget (#49376)

Backport note: rc/1.108.0 also needs graceful_stop_seconds imported in test_logging_only_scope_chaos.py, which main already imported

(cherry picked from commit c8c0248ac739cd4bef2fe27ff635885c01471433)
---
 .../observability/test_arize_otel_v2_openinference_modes.py | 3 ++-
 .../observability/test_logging_only_scope_chaos.py          | 6 +++---
 .../observability/test_s3_v2_partition_granularity.py       | 4 ++--
 3 files changed, 7 insertions(+), 6 deletions(-)

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_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: