From f8cffc2a1e0f4bea4bb7539996a890230eb584bb Mon Sep 17 00:00:00 2001
From: Brandon Allard <brandon@redpanda.com>
Date: Thu, 26 Oct 2023 18:31:21 -0400
Subject: [PATCH] tests: add tests for the wait_ms parameter in the cpu
 profiler API

---
 .../tests/cpu_profiler_admin_api_test.py      | 42 ++++++++++++++++++-
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/tests/rptest/tests/cpu_profiler_admin_api_test.py b/tests/rptest/tests/cpu_profiler_admin_api_test.py
index 956a7cd36525d..38fd91e6d6e03 100644
--- a/tests/rptest/tests/cpu_profiler_admin_api_test.py
+++ b/tests/rptest/tests/cpu_profiler_admin_api_test.py
@@ -15,6 +15,8 @@
 from rptest.services.kgo_repeater_service import repeater_traffic
 from rptest.utils.mode_checks import skip_debug_mode
 
+import requests
+
 
 class CPUProfilerAdminAPITest(RedpandaTest):
     topics = (TopicSpec(partition_count=30, replication_factor=3), )
@@ -26,7 +28,7 @@ def __init__(self, test_context):
             log_config=LoggingConfig('info',
                                      logger_levels={'resources': 'trace'}),
             extra_rp_conf={
-                "cpu_profiler_enabled": True,
+                "cpu_profiler_enabled": False,
                 "cpu_profiler_sample_period_ms": 50,
             })
 
@@ -34,6 +36,10 @@ def __init__(self, test_context):
 
     @cluster(num_nodes=4)
     def test_get_cpu_profile(self):
+        self.redpanda.set_cluster_config({
+            "cpu_profiler_enabled": True,
+        })
+
         # Provide traffic so there is something to sample.
         with repeater_traffic(context=self.test_context,
                               redpanda=self.redpanda,
@@ -41,7 +47,7 @@ def test_get_cpu_profile(self):
                               msg_size=4096,
                               workers=1) as repeater:
             repeater.await_group_ready()
-            repeater.await_progress(1024,
+            repeater.await_progress(2 * 1024,
                                     timeout_sec=100 if self.debug_mode else 75)
 
             profile = self.admin.get_cpu_profile()
@@ -50,3 +56,35 @@ def test_get_cpu_profile(self):
             assert len(
                 profile[0]["samples"]
             ) > 0, "At least one cpu profile should've been collected."
+
+    @cluster(num_nodes=4)
+    def test_get_cpu_profile_with_override(self):
+        # Provide traffic so there is something to sample.
+        with repeater_traffic(context=self.test_context,
+                              redpanda=self.redpanda,
+                              topic=self.topic,
+                              msg_size=4096,
+                              workers=1) as repeater:
+            repeater.await_group_ready()
+            profile = self.admin.get_cpu_profile(wait_ms=30 * 1_000)
+
+            assert len(profile) > 0, "At least one shard should exist"
+            assert len(
+                profile[0]["samples"]
+            ) > 0, "At least one cpu profile should've been collected."
+
+    @cluster(num_nodes=3)
+    def test_get_cpu_profile_with_override_limits(self):
+        try:
+            self.admin.get_cpu_profile(wait_ms=16 * 60 * 1_000)
+        except requests.exceptions.HTTPError:
+            pass
+        else:
+            assert False, "call with wait_ms > 15min should of failed"
+
+        try:
+            self.admin.get_cpu_profile(wait_ms="0")
+        except requests.exceptions.HTTPError:
+            pass
+        else:
+            assert False, "call with wait_ms < 1ms should of failed"