update225822582258$

Signed-off-by: Sai Medhini Reddy Maryada <[email protected]>
opensearch-project · Dec 27, 2023 · 1a01963 · 1a01963
1 parent 17f8eb6
commit 1a01963
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/osbenchmark/client.py b/osbenchmark/client.py
@@ -362,7 +362,15 @@ async def on_request_end(session, trace_config_ctx, params):
         self.client_options["trace_config"] = trace_config
 
         class BenchmarkAsyncOpenSearch(opensearchpy.AsyncOpenSearch, RequestContextHolder):
-            pass
+            print("benchmark###################")
+            async def __aenter__(self) -> Any:
+                print("enter###############")
+                return self
+
+            async def __aexit__(self, *_: Any) -> None:
+                print("exit###############")
+                await self.close()
+            #pass
 
         if "amazon_aws_log_in" not in self.client_options:
             return BenchmarkAsyncOpenSearch(hosts=self.hosts,