tests: support wait_ms in admin class

ballard26 · Nov 2, 2023 · 7032a4c · 7032a4c
1 parent 4b7ea5a
commit 7032a4c
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/tests/rptest/services/admin.py b/tests/rptest/services/admin.py
@@ -1070,11 +1070,23 @@ def get_sampled_memory_profile(self, node=None, shard=None):
                              node=node,
                              **kwargs).json()
 
-    def get_cpu_profile(self, node=None):
+    def get_cpu_profile(self, node=None, wait_ms=None):
         """
         Get the CPU profile of a node.
         """
-        return self._request("get", "debug/cpu_profile", node=node).json()
+        path = "debug/cpu_profile"
+        params = {}
+        timeout = DEFAULT_TIMEOUT
+
+        if wait_ms:
+            params["wait_ms"] = wait_ms
+            timeout = max(2 * (int(wait_ms) // 1_000), timeout)
+
+        return self._request("get",
+                             path,
+                             node=node,
+                             timeout=timeout,
+                             params=params).json()
 
     def get_local_offsets_translated(self,
                                      offsets,