fixes in poll_query_until()

Author: funbringer

testgres/testgres.py

@@ -996,6 +996,7 @@ def poll_query_until(self,
             max_attempts: how many times should we try? 0 == infinite
             sleep_time: how much should we sleep after a failure?
             expected: what should be returned to break the cycle?
+            commit: should (possible) changes be committed?
             raise_programming_error: mute ProgrammingError?
             raise_internal_error: mute InternalError?
         """
@@ -1024,7 +1025,7 @@ def poll_query_until(self,
                 if len(res[0]) == 0:
                     raise QueryException('Query returned 0 columns')
 
-                if res[0][0]:
+                if res[0][0] == expected:
                     return  # done
 
             except pglib.ProgrammingError as e:

tests/test_simple.py

@@ -376,6 +376,7 @@ def test_poll_query_until(self):
             with self.assertRaises(QueryException):
                 node.poll_query_until(
                     dbname='postgres', query='select from pg_class limit 1')
+
             # check None, fail
             with self.assertRaises(QueryException):
                 node.poll_query_until(
@@ -386,17 +387,26 @@ def test_poll_query_until(self):
                 dbname='postgres', query='create table def()',
                 expected=None)    # returns nothing
 
-            # check arbitrary expected value
+            # check arbitrary expected value, fail
+            with self.assertRaises(TimeoutException):
+                node.poll_query_until(
+                    dbname='postgres',
+                    query='select 3',
+                    expected=1,
+                    max_attempts=3,
+                    sleep_time=0.01)
+
+            # check arbitrary expected value, ok
             node.poll_query_until(
-                dbname='postgres', query='select 1', expected=1)
+                dbname='postgres', query='select 2', expected=2)
 
             # check timeout
             with self.assertRaises(TimeoutException):
                 node.poll_query_until(
                     dbname='postgres',
                     query='select 1 > 2',
-                    max_attempts=5,
-                    sleep_time=0.2)
+                    max_attempts=3,
+                    sleep_time=0.01)
 
     def test_logging(self):
         logfile = tempfile.NamedTemporaryFile('w', delete=True)