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)