From aac13d136e6f08587b8cca37e5c9a3c463ec3d01 Mon Sep 17 00:00:00 2001
From: Brandon Allard <brandon@redpanda.com>
Date: Thu, 26 Oct 2023 18:30:38 -0400
Subject: [PATCH] tests: support wait_ms in admin class

---
 tests/rptest/services/admin.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/tests/rptest/services/admin.py b/tests/rptest/services/admin.py
index 09954267a8390..484a9dfff23c7 100644
--- 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,