From a68dec473591f2af8b21bc8117f22dc49ed25f89 Mon Sep 17 00:00:00 2001
From: vshepard <v.shepard@postgrespro.ru>
Date: Tue, 21 May 2024 00:22:23 +0200
Subject: [PATCH] Add pgbench_with_wait function

---
 testgres/node.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/testgres/node.py b/testgres/node.py
index c9b34d7b..d1784cb9 100644
--- 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().