Add request start and end parameter to sleep runner (#491)

Signed-off-by: Vijayan Balasubramanian <[email protected]>

osbenchmark/worker_coordinator/runner.py

@@ -1675,10 +1675,13 @@ class Sleep(Runner):
     """
 
     async def __call__(self, opensearch, params):
+        sleep_duration = mandatory(params, "duration", "sleep")
         opensearch.on_request_start()
         try:
-            await asyncio.sleep(mandatory(params, "duration", "sleep"))
+            request_context_holder.on_client_request_start()
+            await asyncio.sleep(sleep_duration)
         finally:
+            request_context_holder.on_client_request_end()
             opensearch.on_request_end()
 
     def __repr__(self, *args, **kwargs):

tests/worker_coordinator/runner_test.py

@@ -3927,32 +3927,40 @@ async def test_raw_with_timeout_and_opaqueid(self, opensearch, on_client_request
 
 class SleepTests(TestCase):
     @mock.patch("opensearchpy.OpenSearch")
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_end')
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_start')
     # To avoid real sleeps in unit tests
     @mock.patch("asyncio.sleep", return_value=as_future())
     @run_async
-    async def test_missing_parameter(self, sleep, opensearch):
+    async def test_missing_parameter(self, sleep, on_client_request_start, on_client_request_end, opensearch):
         r = runner.Sleep()
         with self.assertRaisesRegex(exceptions.DataError,
                                     "Parameter source for operation 'sleep' did not provide the mandatory parameter "
                                     "'duration'. Add it to your parameter source and try again."):
             await r(opensearch, params={})
 
         self.assertEqual(0, opensearch.call_count)
-        self.assertEqual(1, opensearch.on_request_start.call_count)
-        self.assertEqual(1, opensearch.on_request_end.call_count)
+        self.assertEqual(0, opensearch.on_request_start.call_count)
+        self.assertEqual(0, opensearch.on_request_end.call_count)
+        self.assertEqual(0, on_client_request_start.call_count)
+        self.assertEqual(0, on_client_request_end.call_count)
         self.assertEqual(0, sleep.call_count)
 
     @mock.patch("opensearchpy.OpenSearch")
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_end')
+    @mock.patch('osbenchmark.client.RequestContextHolder.on_client_request_start')
     # To avoid real sleeps in unit tests
     @mock.patch("asyncio.sleep", return_value=as_future())
     @run_async
-    async def test_sleep(self, sleep, opensearch):
+    async def test_sleep(self, sleep, on_client_request_start, on_client_request_end, opensearch):
         r = runner.Sleep()
         await r(opensearch, params={"duration": 4.3})
 
         self.assertEqual(0, opensearch.call_count)
         self.assertEqual(1, opensearch.on_request_start.call_count)
         self.assertEqual(1, opensearch.on_request_end.call_count)
+        self.assertEqual(1, on_client_request_start.call_count)
+        self.assertEqual(1, on_client_request_end.call_count)
         sleep.assert_called_once_with(4.3)