added support for multiple users

Author: zilder

setup.py

@@ -2,7 +2,7 @@
 setup(
     name='testgres',
     packages=['testgres'],
-    version='0.1.13',
+    version='0.1.14',
     description='Testing utility for postgresql and its extensions',
     author='Ildar Musin',
     author_email='zildermann@gmail.com',

testgres/testgres.py

@@ -211,7 +211,7 @@ def init(self, allows_streaming=False):
             conf.write(
                 # "unix_socket_directories = '%s'\n"
                 # "listen_addresses = ''\n";)
-                "listen_addresses = {}\n".format(self.host))
+                "listen_addresses = '{}'\n".format(self.host))
 
             if allows_streaming:
                 # TODO: wal_level = hot_standby (9.5)
@@ -355,7 +355,7 @@ def cleanup(self):
 
         return self
 
-    def psql(self, dbname, query=None, filename=None):
+    def psql(self, dbname, query=None, filename=None, username=None):
         """Executes a query by the psql
 
         Returns a tuple (code, stdout, stderr) in which:
@@ -375,6 +375,10 @@ def psql(self, dbname, query=None, filename=None):
         else:
             raise QueryException('Query or filename must be provided')
 
+        # Specify user if needed
+        if username:
+            psql_params.extend(("-U", username))
+
         # start psql process
         process = subprocess.Popen(
             psql_params,
@@ -386,13 +390,13 @@ def psql(self, dbname, query=None, filename=None):
         out, err = process.communicate()
         return process.returncode, out, err
 
-    def safe_psql(self, dbname, query):
+    def safe_psql(self, dbname, query, username=None):
         """Executes a query by the psql
 
         Returns the stdout if succeed. Otherwise throws the
         ClusterException with stderr output
         """
-        ret, out, err = self.psql(dbname, query)
+        ret, out, err = self.psql(dbname, query, username=username)
         if ret:
             raise ClusterException("psql failed:\n" + six.text_type(err))
         return out
@@ -441,9 +445,9 @@ def poll_query_until(self, dbname, query):
             attemps += 1
         raise QueryException("Timeout while waiting for query to return True")
 
-    def execute(self, dbname, query):
+    def execute(self, dbname, query, username=None):
         """Executes the query and returns all rows"""
-        with self.connect(dbname) as node_con:
+        with self.connect(dbname, username) as node_con:
             return node_con.execute(query)
 
     def backup(self, name):
@@ -465,8 +469,8 @@ def backup(self, name):
 
             return backup_path
 
-    def connect(self, dbname='postgres'):
-        return NodeConnection(parent_node=self, dbname=dbname)
+    def connect(self, dbname='postgres', username=None):
+        return NodeConnection(parent_node=self, dbname=dbname, user=username)
 
 
 def get_username():

testgres/tests/test_simple.py

@@ -1,74 +1,80 @@
 import unittest
-# import time
-# from testgres import clean_all
 from testgres import get_new_node, stop_all
 
 
 class SimpleTest(unittest.TestCase):
 
-	def teardown(self):
-		# clean_all()
-		stop_all()
+    def teardown(self):
+        # clean_all()
+        stop_all()
 
-	@unittest.skip("demo")
-	def test_start_stop(self):
-		node = get_new_node('test')
-		node.init()
-		node.start()
-		res = node.execute('postgres', 'select 1')
-		self.assertEqual(len(res), 1)
-		self.assertEqual(res[0][0], 1)
-		node.stop()
+    @unittest.skip("demo")
+    def test_start_stop(self):
+        node = get_new_node('test')
+        node.init()
+        node.start()
+        res = node.execute('postgres', 'select 1')
+        self.assertEqual(len(res), 1)
+        self.assertEqual(res[0][0], 1)
+        node.stop()
 
-	def test_backup_and_replication(self):
-		node = get_new_node('test')
-		replica = get_new_node('repl')
+    def test_backup_and_replication(self):
+        node = get_new_node('test')
+        replica = get_new_node('repl')
 
-		node.init(allows_streaming=True)
-		node.start()
-		node.psql('postgres', 'create table abc(a int, b int)')
-		node.psql('postgres', 'insert into abc values (1, 2)')
-		node.backup('my_backup')
+        node.init(allows_streaming=True)
+        node.start()
+        node.psql('postgres', 'create table abc(a int, b int)')
+        node.psql('postgres', 'insert into abc values (1, 2)')
+        node.backup('my_backup')
 
-		replica.init_from_backup(node, 'my_backup', has_streaming=True)
-		replica.start()
-		res = replica.execute('postgres', 'select * from abc')
-		self.assertEqual(len(res), 1)
-		self.assertEqual(res[0], (1, 2))
+        replica.init_from_backup(node, 'my_backup', has_streaming=True)
+        replica.start()
+        res = replica.execute('postgres', 'select * from abc')
+        self.assertEqual(len(res), 1)
+        self.assertEqual(res[0], (1, 2))
 
-		# Insert into master node
-		node.psql('postgres', 'insert into abc values (3, 4)')
-		# Wait until data syncronizes
-		node.poll_query_until(
-			'postgres',
-			'SELECT pg_current_xlog_location() <= replay_location '
-			'FROM pg_stat_replication WHERE application_name = \'%s\''
-			% replica.name)
-		# time.sleep(0.5)
-		# Check that this record was exported to replica
-		res = replica.execute('postgres', 'select * from abc')
-		self.assertEqual(len(res), 2)
-		self.assertEqual(res[1], (3, 4))
+        # Insert into master node
+        node.psql('postgres', 'insert into abc values (3, 4)')
+        # Wait until data syncronizes
+        node.poll_query_until(
+            'postgres',
+            'SELECT pg_current_xlog_location() <= replay_location '
+            'FROM pg_stat_replication WHERE application_name = \'%s\''
+            % replica.name)
+        # time.sleep(0.5)
+        # Check that this record was exported to replica
+        res = replica.execute('postgres', 'select * from abc')
+        self.assertEqual(len(res), 2)
+        self.assertEqual(res[1], (3, 4))
 
-		node.stop()
-		replica.stop()
+        node.stop()
+        replica.stop()
+
+    def test_dump(self):
+        node = get_new_node('test')
+        node.init().start()
+        node.safe_psql(
+            'postgres',
+            'create table abc as '
+            'select g as a, g as b from generate_series(1, 10) as g'
+        )
+        node.psql('postgres', 'create database test')
+        node.dump('postgres', 'test.sql')
+        node.restore('test', 'test.sql')
+        self.assertEqual(
+            node.psql('postgres', 'select * from abc'),
+            node.psql('test', 'select * from abc'),
+        )
+        node.stop()
+
+    def test_users(self):
+        node = get_new_node('master')
+        node.init().start()
+        node.psql('postgres', 'create role test_user login')
+        value = node.safe_psql('postgres', 'select 1', username='test_user')
+        self.assertEqual(value, '1\n')
 
-	def test_dump(self):
-		node = get_new_node('test')
-		node.init().start()
-		node.safe_psql(
-			'postgres',
-			'create table abc as '
-			'select g as a, g as b from generate_series(1, 10) as g'
-		)
-		node.psql('postgres', 'create database test')
-		node.dump('postgres', 'test.sql')
-		node.restore('test', 'test.sql')
-		self.assertEqual(
-			node.psql('postgres', 'select * from abc'),
-			node.psql('test', 'select * from abc'),
-		)
-		node.stop()
 
 if __name__ == '__main__':
-	unittest.main()
+    unittest.main()