Merge pull request #51 from funbringer/poll_query_until_suppress

funbringer · web-flow · commit 0847cf2ae3a9 · 2018-07-11T14:12:21.000+03:00
add 'suppress' to PostgresNode.poll_query_until() to mute exceptions
diff --git a/testgres/__init__.py b/testgres/__init__.py
@@ -10,8 +10,10 @@
 
 from .connection import \
     NodeConnection, \
+    DatabaseError, \
     InternalError, \
-    ProgrammingError
+    ProgrammingError, \
+    OperationalError
 
 from .exceptions import *
 from .enums import *
diff --git a/testgres/connection.py b/testgres/connection.py
@@ -17,9 +17,11 @@
 
 from .exceptions import QueryException
 
-# export these exceptions
+# export some exceptions
+DatabaseError = pglib.DatabaseError
 InternalError = pglib.InternalError
 ProgrammingError = pglib.ProgrammingError
+OperationalError = pglib.OperationalError
 
 
 class NodeConnection(object):
diff --git a/testgres/node.py b/testgres/node.py
@@ -19,10 +19,7 @@
 
 from .config import testgres_config
 
-from .connection import \
-    NodeConnection, \
-    InternalError,  \
-    ProgrammingError
+from .connection import NodeConnection
 
 from .consts import \
     DATA_DIR, \
@@ -541,10 +538,10 @@ def append_conf(self, line='', filename=PG_CONF_FILE, **kwargs):
             This instance of :class:`.PostgresNode`.
 
         Examples:
-            append_conf(fsync=False)
-            append_conf('log_connections = yes')
-            append_conf(random_page_cost=1.5, fsync=True, ...)
-            append_conf('postgresql.conf', 'synchronous_commit = off')
+            >>> append_conf(fsync=False)
+            >>> append_conf('log_connections = yes')
+            >>> append_conf(random_page_cost=1.5, fsync=True, ...)
+            >>> append_conf('postgresql.conf', 'synchronous_commit = off')
         """
 
         lines = [line]
@@ -980,8 +977,7 @@ def poll_query_until(self,
                          sleep_time=1,
                          expected=True,
                          commit=True,
-                         raise_programming_error=True,
-                         raise_internal_error=True):
+                         suppress=None):
         """
         Run a query once per second until it returns 'expected'.
         Query should return a single value (1 row, 1 column).
@@ -994,13 +990,13 @@ def poll_query_until(self,
             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: enable ProgrammingError?
-            raise_internal_error: enable InternalError?
+            suppress: a collection of exceptions to be suppressed.
 
         Examples:
-            poll_query_until('select true')
-            poll_query_until('postgres', "select now() > '01.01.2018'")
-            poll_query_until('select false', expected=True, max_attempts=4)
+            >>> poll_query_until('select true')
+            >>> poll_query_until('postgres', "select now() > '01.01.2018'")
+            >>> poll_query_until('select false', expected=True, max_attempts=4)
+            >>> poll_query_until('select 1', suppress={testgres.OperationalError})
         """
 
         # sanity checks
@@ -1032,13 +1028,8 @@ def poll_query_until(self,
                 elif expected is None:
                     return    # done
 
-            except ProgrammingError as e:
-                if raise_programming_error:
-                    raise e
-
-            except InternalError as e:
-                if raise_internal_error:
-                    raise e
+            except tuple(suppress or []):
+                pass    # we're suppressing them
 
             time.sleep(sleep_time)
             attempts += 1
@@ -1229,13 +1220,14 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
             options: additional options for pgbench (list).
 
             **kwargs: named options for pgbench.
-                Examples:
-                    pgbench_run(initialize=True, scale=2)
-                    pgbench_run(time=10)
                 Run pgbench --help to learn more.
 
         Returns:
             Stdout produced by pgbench.
+
+        Examples:
+            >>> pgbench_run(initialize=True, scale=2)
+            >>> pgbench_run(time=10)
         """
 
         # Set default arguments
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -626,7 +626,7 @@ def test_poll_query_until(self):
                     query='dummy2',
                     max_attempts=3,
                     sleep_time=0.01,
-                    raise_programming_error=False)
+                    suppress={testgres.ProgrammingError})
 
             # check 1 arg, ok
             node.poll_query_until('select true')
