improve poll_query_until()

Author: funbringer

testgres/testgres.py

@@ -93,6 +93,10 @@ class QueryException(Exception):
     pass
 
 
+class TimeoutException(Exception):
+    pass
+
+
 class StartNodeException(Exception):
     pass
 
@@ -115,8 +119,8 @@ def __init__(self, node_name, fd):
 
         threading.Thread.__init__(self)
 
-        self.node_name = node_name
         self.fd = fd
+        self.node_name = node_name
         self.stop_event = threading.Event()
         self.logger = logging.getLogger(node_name)
         self.logger.setLevel(logging.INFO)
@@ -366,13 +370,19 @@ def __bool__(self):
 
 
 class PostgresNode(object):
-    def __init__(self, name, port=None, base_dir=None, use_logging=False):
+    def __init__(self,
+                 name,
+                 port=None,
+                 base_dir=None,
+                 use_logging=False,
+                 master=None):
         global bound_ports
 
         # check that port is not used
         if port in bound_ports:
             raise InitNodeException('port {} is already in use'.format(port))
 
+        self.master = master
         self.name = name
         self.host = '127.0.0.1'
         self.port = port or reserve_port()
@@ -809,13 +819,22 @@ def poll_query_until(self, dbname, query, username=None, max_attempts=60, sleep_
                                username=username,
                                commit=True)
 
+            if res is None:
+                raise QueryException('Query returned None')
+
+            if len(res) == 0:
+                raise QueryException('Query returned 0 rows')
+
+            if len(res[0]) == 0:
+                raise QueryException('Query returned 0 columns')
+
             if res[0][0]:
                 return  # done
 
             time.sleep(sleep_time)
             attemps += 1
 
-        raise QueryException('Query timeout')
+        raise TimeoutException('Query timeout')
 
     def execute(self, dbname, query, username=None, commit=False):
         """

tests/test_simple.py

@@ -9,7 +9,10 @@
 
 from distutils.version import LooseVersion
 
-from testgres import InitNodeException, StartNodeException, ExecUtilException, BackupException
+from testgres import InitNodeException, \
+    StartNodeException, ExecUtilException, \
+    BackupException, QueryException
+
 from testgres import get_new_node, get_pg_config
 from testgres import bound_ports
 from testgres import NodeStatus
@@ -283,6 +286,30 @@ def test_poll_query_until(self):
 
             self.assertTrue(end_time - start_time >= 5)
 
+            # check 0 rows
+            got_exception = False
+            try:
+                node.poll_query_until('postgres', 'select * from pg_class where true = false')
+            except QueryException as e:
+                got_exception = True
+            self.assertTrue(got_exception)
+
+            # check 0 columns
+            got_exception = False
+            try:
+                node.poll_query_until('postgres', 'select from pg_class limit 1')
+            except QueryException as e:
+                got_exception = True
+            self.assertTrue(got_exception)
+
+            # check None
+            got_exception = False
+            try:
+                node.poll_query_until('postgres', 'create table abc (val int)')
+            except QueryException as e:
+                got_exception = True
+            self.assertTrue(got_exception)
+
     def test_logging(self):
         regex = re.compile('.+?LOG:.*')
         logfile = tempfile.NamedTemporaryFile('w', delete=True)