Add pgbench_with_wait function (#122)

demonolock · web-flow · commit 8a25cb375138 · 2024-05-21T01:57:29.000+03:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -1371,7 +1371,7 @@ def pgbench(self,
                 username=None,
                 stdout=None,
                 stderr=None,
-                options=[]):
+                options=None):
         """
         Spawn a pgbench process.
 
@@ -1385,6 +1385,8 @@ def pgbench(self,
         Returns:
             Process created by subprocess.Popen.
         """
+        if options is None:
+            options = []
 
         # Set default arguments
         dbname = dbname or default_dbname()
@@ -1404,6 +1406,29 @@ def pgbench(self,
 
         return proc
 
+    def pgbench_with_wait(self,
+                          dbname=None,
+                          username=None,
+                          stdout=None,
+                          stderr=None,
+                          options=None):
+        """
+        Do pgbench command and wait.
+
+        Args:
+            dbname: database name to connect to.
+            username: database user name.
+            stdout: stdout file to be used by Popen.
+            stderr: stderr file to be used by Popen.
+            options: additional options for pgbench (list).
+        """
+        if options is None:
+            options = []
+
+        with self.pgbench(dbname, username, stdout, stderr, options) as pgbench:
+            pgbench.wait()
+        return
+
     def pgbench_init(self, **kwargs):
         """
         Small wrapper for pgbench_run().
