add username arg where necessary, more tests

Author: funbringer

testgres/testgres.py

@@ -164,17 +164,17 @@ def __init__(self,
                  parent_node,
                  dbname,
                  host="127.0.0.1",
-                 user=None,
+                 username=None,
                  password=None):
 
         # Use default user if not specified
-        user = user or _default_username()
+        username = username or _default_username()
 
         self.parent_node = parent_node
 
         self.connection = pglib.connect(
             database=dbname,
-            user=user,
+            user=username,
             port=parent_node.port,
             host=host,
             password=password)
@@ -251,10 +251,17 @@ class NodeBackup(object):
     def log_file(self):
         return os.path.join(self.base_dir, BACKUP_LOG_FILE)
 
-    def __init__(self, node, base_dir=None, xlog_method=DEFAULT_XLOG_METHOD):
+    def __init__(self,
+                 node,
+                 base_dir=None,
+                 username=None,
+                 xlog_method=DEFAULT_XLOG_METHOD):
+
         if not node.status():
             raise BackupException('Node must be running')
 
+        # set default arguments
+        username = username or _default_username()
         base_dir = base_dir or tempfile.mkdtemp()
 
         # create directory if needed
@@ -269,7 +276,8 @@ def __init__(self, node, base_dir=None, xlog_method=DEFAULT_XLOG_METHOD):
         _params = [
             "-D{}".format(data_dir),
             "-p{}".format(node.port),
-            "-X", xlog_method
+            "-U{}".format(username),
+            "-X{}".format(xlog_method)
         ]
         _execute_utility("pg_basebackup", _params, self.log_file)
 
@@ -782,7 +790,7 @@ def dump(self, dbname, filename=None):
 
         return filename
 
-    def restore(self, dbname, filename):
+    def restore(self, dbname, filename, username=None):
         """
         Restore database from pg_dump's file.
 
@@ -791,9 +799,9 @@ def restore(self, dbname, filename):
             filename: database dump taken by pg_dump (str).
         """
 
-        self.psql(dbname, filename=filename)
+        self.psql(dbname=dbname, filename=filename, username=username)
 
-    def poll_query_until(self, dbname, query):
+    def poll_query_until(self, dbname, query, username=None, max_attempts=60, sleep_time=1):
         """
         Run a query once a second until it returs True.
 
@@ -802,16 +810,17 @@ def poll_query_until(self, dbname, query):
             query: query to be executed (str).
         """
 
-        max_attemps = 60
         attemps = 0
+        while attemps < max_attempts:
+            res = self.execute(dbname=dbname,
+                               query=query,
+                               username=username,
+                               commit=True)
 
-        while attemps < max_attemps:
-            ret = self.safe_psql(dbname, query)
-            # TODO: fix psql so that it returns result without newline
-            if ret == six.b("t\n"):
-                return
+            if res[0][0]:
+                return  # done
 
-            time.sleep(1)
+            time.sleep(sleep_time)
             attemps += 1
 
         raise QueryException('Query timeout')
@@ -836,19 +845,21 @@ def execute(self, dbname, query, username=None, commit=False):
                 node_con.commit()
             return res
 
-    def backup(self, xlog_method=DEFAULT_XLOG_METHOD):
+    def backup(self, username=None, xlog_method=DEFAULT_XLOG_METHOD):
         """
         Perform pg_basebackup.
 
         Args:
-            xlog_method: a method for collecting the logs ('fetch' | 'stream')
+            username: database user name (str).
+            xlog_method: a method for collecting the logs ('fetch' | 'stream').
 
         Returns:
             A smart object of type NodeBackup.
         """
 
-        # NodeBackup will handle this
-        return NodeBackup(self, xlog_method=xlog_method)
+        return NodeBackup(node=self,
+                          username=username,
+                          xlog_method=xlog_method)
 
     def pgbench_init(self, dbname='postgres', scale=1, options=[]):
         """
@@ -905,7 +916,9 @@ def connect(self, dbname='postgres', username=None):
             An instance of NodeConnection.
         """
 
-        return NodeConnection(parent_node=self, dbname=dbname, user=username)
+        return NodeConnection(parent_node=self,
+                              dbname=dbname,
+                              username=username)
 
 
 def _default_username():

testgres/tests/test_simple.py

@@ -160,7 +160,7 @@ def test_backup_simple(self):
             master.psql('postgres',
                         'create table test as select generate_series(1, 4) i')
 
-            with master.backup('stream') as backup:
+            with master.backup(xlog_method='stream') as backup:
                 with backup.spawn_primary('slave') as slave:
                     slave.start()
                     res = slave.execute('postgres',
@@ -171,12 +171,12 @@ def test_backup_multiple(self):
         with get_new_node('node') as node:
             node.init(allow_streaming=True).start()
 
-            with node.backup('fetch') as backup1, \
-                    node.backup('fetch') as backup2:
+            with node.backup(xlog_method='fetch') as backup1, \
+                    node.backup(xlog_method='fetch') as backup2:
 
                 self.assertNotEqual(backup1.base_dir, backup2.base_dir)
 
-            with node.backup('fetch') as backup:
+            with node.backup(xlog_method='fetch') as backup:
                 with backup.spawn_primary('node1', destroy=False) as node1, \
                         backup.spawn_primary('node2', destroy=False) as node2:
 
@@ -186,7 +186,7 @@ def test_backup_exhaust(self):
         with get_new_node('node') as node:
             node.init(allow_streaming=True).start()
 
-            with node.backup('fetch') as backup:
+            with node.backup(xlog_method='fetch') as backup:
                 with backup.spawn_primary('node1') as node1:
                     pass
 
@@ -264,6 +264,19 @@ def test_users(self):
             value = node.safe_psql('postgres', 'select 1', username='test_user')
             self.assertEqual(value, six.b('1\n'))
 
+    def test_poll_query_until(self):
+        with get_new_node('master') as node:
+            node.init().start()
+
+            get_time = 'select extract(epoch from now())'
+            check_time = 'select extract(epoch from now()) - {} >= 5'
+
+            start_time = node.execute('postgres', get_time)[0][0]
+            node.poll_query_until('postgres', check_time.format(start_time))
+            end_time = node.execute('postgres', get_time)[0][0]
+
+            self.assertTrue(end_time - start_time >= 5)
+
     def test_logging(self):
         regex = re.compile('.+?LOG:.*')
         logfile = tempfile.NamedTemporaryFile('w', delete=True)