From 603992f658262a201cf9546788915543dea08453 Mon Sep 17 00:00:00 2001
From: Eugene Ostroukhov <eostroukhov@google.com>
Date: Thu, 9 Jan 2025 11:11:35 -0800
Subject: [PATCH] [fallback] Make the `get_stats` timeout configurable and
 default to 5m

Go test seems to fail because timeout is too short. Making it
configurable should help with the debugging. Also, setting it
to 5m may just fix the issue.
---
 framework/helpers/docker.py | 4 +++-
 tests/fallback_test.py      | 7 ++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/framework/helpers/docker.py b/framework/helpers/docker.py
index c83bc5be..500b70d3 100644
--- a/framework/helpers/docker.py
+++ b/framework/helpers/docker.py
@@ -260,6 +260,7 @@ def __init__(
         name: str,
         url: str,
         image: str,
+        stats_request_timeout_s: str,
     ):
         super().__init__(
             manager=manager,
@@ -278,13 +279,14 @@ def __init__(
                 }
             },
         )
+        self.stats_request_timeout_s = stats_request_timeout_s
 
     def get_stats(self, num_rpcs: int):
         logger.debug("Sending %d requests", num_rpcs)
         stub = test_pb2_grpc.LoadBalancerStatsServiceStub(self.channel())
         res = stub.GetClientStats(
             messages_pb2.LoadBalancerStatsRequest(
-                num_rpcs=num_rpcs, timeout_sec=math.ceil(num_rpcs * 10)
+                num_rpcs=num_rpcs, timeout_sec=self.stats_request_timeout_s
             )
         )
         return res
diff --git a/tests/fallback_test.py b/tests/fallback_test.py
index a9618391..ed979107 100644
--- a/tests/fallback_test.py
+++ b/tests/fallback_test.py
@@ -49,7 +49,11 @@
 _STATUS_POLL_INTERVAL_MS = flags.DEFINE_integer(
     "status_poll_interval_ms", 300, "Channel status poll interval (in ms)"
 )
-
+_STATS_REQUEST_TIMEOUT_S = flags.DEFINE_integer(
+    "stats_request_timeout_s",
+    300,
+    "Number of seconds the client will wait for the requested number of RPCs",
+)
 _LISTENER = "listener_0"
 
 absl.flags.adopt_module_key_flags(framework.xds_k8s_testcase)
@@ -92,6 +96,7 @@ def start_client(self, port: int = None, name: str = None):
             port=port or get_free_port(),
             url=f"xds:///{_LISTENER}",
             image=framework.xds_k8s_flags.CLIENT_IMAGE.value,
+            stats_request_timeout_s = _STATS_REQUEST_TIMEOUT_S.value,
         )
 
     def start_control_plane(