Merge pull request #15 from stalkerg/master

funbringer · web-flow · commit df65a94dc40d · 2016-12-12T19:01:51.000+03:00
Add pgbench init and pgbench command for node.
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 setup(
     name='testgres',
     packages=['testgres'],
-    version='0.1.16',
+    version='0.1.17',
     description='Testing utility for postgresql and its extensions',
     author='Ildar Musin',
     author_email='zildermann@gmail.com',
diff --git a/testgres/testgres.py b/testgres/testgres.py
@@ -472,6 +472,42 @@ def backup(self, name):
 
             return backup_path
 
+    def pgbench_init(self, dbname='postgres', scale=1, options=[]):
+        """Prepare pgbench database"""
+        pgbench = self.get_bin_path("pgbench")
+        params = [
+            pgbench,
+            "-i",
+            "-s", "%i" % scale,
+            "-p", "%i" % self.port
+        ] + options + [dbname]
+        with open(self.output_filename, "a") as file_out, \
+                open(self.error_filename, "a") as file_err:
+            ret = subprocess.call(
+                params,
+                stdout=file_out,
+                stderr=file_err
+            )
+            if ret:
+                raise ClusterException("pgbench init failed")
+
+            return True
+
+    def pgbench(self, dbname='postgres', stdout=None, stderr=None, options=[]):
+        """Make pgbench process"""
+        pgbench = self.get_bin_path("pgbench")
+        params = [
+            pgbench,
+            "-p", "%i" % self.port
+        ] + options + [dbname]
+        proc = subprocess.Popen(
+            params,
+            stdout=stdout,
+            stderr=stderr
+        )
+
+        return proc
+
     def connect(self, dbname='postgres', username=None):
         return NodeConnection(parent_node=self, dbname=dbname, user=username)
 
