fix: fix several issues in eventually_consistent

gcp_devrel/testing/eventually_consistent.py

@@ -18,7 +18,7 @@
 from google.cloud import exceptions
 from retrying import retry
 
-WAIT_EXPONENTIAL_MULTIPLIER = 1000
+WAIT_EXPONENTIAL_MULTIPLIER_DEFAULT = 1000
 WAIT_EXPONENTIAL_MAX_DEFAULT = 30000
 STOP_MAX_ATTEMPT_NUMBER_DEFAULT = 10
 
@@ -29,20 +29,49 @@ def inner(e):
             print('Retrying due to eventual consistency.')
             return True
         return False
+    return inner
 
 
-def mark(f):
-    """Marks an entire test as eventually consistent and retries."""
+def mark(*args, **kwargs):
+    """Marks an entire test as eventually consistent and retries.
+
+    Args:
+        tries: The number of retries.
+        exceptions: The exceptions on which it will retry. It can be
+            single value or a tuple.
+        wait_exponential_multiplier: The exponential multiplier in
+            milliseconds.
+        wait_exponential_max: The maximum wait before the next try in
+            milliseconds.
+    """
     __tracebackhide__ = True
-    return retry(
-        wait_exponential_multiplier=WAIT_EXPONENTIAL_MULTIPLIER,
-        wait_exponential_max=WAIT_EXPONENTIAL_MAX_DEFAULT,
-        stop_max_attempt_number=STOP_MAX_ATTEMPT_NUMBER_DEFAULT,
-        retry_on_exception=_retry_on_exception(
-            (AssertionError, exceptions.GoogleCloudError)))(f)
+    tries = kwargs.get('tries', STOP_MAX_ATTEMPT_NUMBER_DEFAULT)
+    retry_exceptions = kwargs.get(
+        'exceptions', (AssertionError, exceptions.GoogleCloudError))
+    wait_exponential_multiplier = kwargs.get(
+        'wait_exponential_multiplier', WAIT_EXPONENTIAL_MULTIPLIER_DEFAULT)
+    wait_exponential_max = kwargs.get(
+        'wait_exponential_max', WAIT_EXPONENTIAL_MAX_DEFAULT)
+    # support both `@mark` and `@mark()` syntax
+    if len(args) == 1 and callable(args[0]):
+        return retry(
+            wait_exponential_multiplier=wait_exponential_multiplier,
+            wait_exponential_max=wait_exponential_max,
+            stop_max_attempt_number=tries,
+            retry_on_exception=_retry_on_exception(retry_exceptions))(args[0])
 
+    # `mark()` syntax
+    def inner(f):
+        __tracebackhide__ = True
+        return retry(
+            wait_exponential_multiplier=wait_exponential_multiplier,
+            wait_exponential_max=wait_exponential_max,
+            stop_max_attempt_number=tries,
+            retry_on_exception=_retry_on_exception(retry_exceptions))(f)
+    return inner
 
-def call(f, exceptions=AssertionError, tries=STOP_MAX_ATTEMPT_NUMBER_DEFAULT):
+
+def call(*args, **kwargs):
     """Call a given function and treat it as eventually consistent.
 
     The function will be called immediately and retried with exponential
@@ -51,17 +80,56 @@ def call(f, exceptions=AssertionError, tries=STOP_MAX_ATTEMPT_NUMBER_DEFAULT):
     By default, it only retries on AssertionErrors, but can be told to retry
     on other errors.
 
-    For example:
+    Args:
+        tries: The number of retries.
+        exceptions: The exceptions on which it will retry. It can be
+            single value or a tuple.
+        wait_exponential_multiplier: The exponential multiplier in
+            milliseconds.
+        wait_exponential_max: The maximum wait before the next try in
+            milliseconds.
+
+    Examples:
 
         @eventually_consistent.call
         def _():
             results = client.query().fetch(10)
             assert len(results) == 10
 
+        @eventually_consistent.call(tries=2)
+        def _():
+            results = client.query().fetch(10)
+            assert len(results) == 10
+
+        @eventually_consistent.call(tries=2, exceptions=SomeException)
+        def _():
+            # It might throw SomeException
+            results = client.query().fetch(10)
+
     """
     __tracebackhide__ = True
-    return retry(
-        wait_exponential_multiplier=WAIT_EXPONENTIAL_MULTIPLIER,
-        wait_exponential_max=WAIT_EXPONENTIAL_MAX_DEFAULT,
-        stop_max_attempt_number=tries,
-        retry_on_exception=_retry_on_exception(exceptions))(f)()
+    tries = kwargs.get('tries', STOP_MAX_ATTEMPT_NUMBER_DEFAULT)
+    retry_exceptions = kwargs.get('exceptions', AssertionError)
+    wait_exponential_multiplier = kwargs.get(
+        'wait_exponential_multiplier', WAIT_EXPONENTIAL_MULTIPLIER_DEFAULT)
+    wait_exponential_max = kwargs.get(
+        'wait_exponential_max', WAIT_EXPONENTIAL_MAX_DEFAULT)
+
+    # support both `@call` and `@call()` syntax
+    if len(args) == 1 and callable(args[0]):
+        return retry(
+            wait_exponential_multiplier=wait_exponential_multiplier,
+            wait_exponential_max=wait_exponential_max,
+            stop_max_attempt_number=tries,
+            retry_on_exception=_retry_on_exception(
+                retry_exceptions))(args[0])()
+
+    # `@call()` syntax
+    def inner(f):
+        __tracebackhide__ = True
+        return retry(
+            wait_exponential_multiplier=wait_exponential_multiplier,
+            wait_exponential_max=wait_exponential_max,
+            stop_max_attempt_number=tries,
+            retry_on_exception=_retry_on_exception(retry_exceptions))(f)()
+    return inner

gcp_devrel/testing/eventually_consistent_test.py

@@ -0,0 +1,136 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from gcp_devrel.testing import eventually_consistent
+
+test_mark_simple_tries = 0
+test_mark_args_tries = 0
+test_mark_custom_exception_tries = 0
+test_mark_custom_exception_with_tuple_tries = 0
+
+
+class MyException(Exception):
+    pass
+
+
+@eventually_consistent.mark
+def test_mark_simple():
+    global test_mark_simple_tries
+    test_mark_simple_tries += 1
+    if test_mark_simple_tries == 2:
+        assert True
+    else:
+        assert False
+
+
+@eventually_consistent.mark(tries=2)
+def test_mark_args():
+    global test_mark_args_tries
+    test_mark_args_tries += 1
+    if test_mark_args_tries == 2:
+        assert True
+    else:
+        assert False
+
+
+@eventually_consistent.mark(tries=2, exceptions=MyException)
+def test_mark_custom_exception():
+    global test_mark_custom_exception_tries
+    test_mark_custom_exception_tries += 1
+    if test_mark_custom_exception_tries == 2:
+        assert True
+    else:
+        raise MyException
+
+
+@eventually_consistent.mark(tries=3, exceptions=(MyException, AssertionError))
+def test_mark_custom_exceptions_with_tuple():
+    global test_mark_custom_exception_with_tuple_tries
+    test_mark_custom_exception_with_tuple_tries += 1
+    if test_mark_custom_exception_with_tuple_tries == 3:
+        assert True
+    elif test_mark_custom_exception_with_tuple_tries == 2:
+        raise MyException
+    else:
+        assert False
+
+
+def test_call_simple():
+    tried = 0
+    @eventually_consistent.call
+    def _():
+        nonlocal tried
+        tried += 1
+        if tried == 2:
+            assert True
+        else:
+            assert False
+
+
+def test_call_args():
+    tried = 0
+
+    @eventually_consistent.call(tries=2)
+    def _():
+        nonlocal tried
+        tried += 1
+        if tried == 2:
+            assert True
+        else:
+            assert False
+
+
+def test_call_args_fail():
+    with pytest.raises(AssertionError):
+        tried = 0
+
+        @eventually_consistent.call(tries=2)
+        def _():
+            nonlocal tried
+            tried += 1
+            if tried == 3:
+                assert True
+            else:
+                assert False
+
+
+def test_call_custom_exception():
+    tried = 0
+
+    @eventually_consistent.call(tries=2, exceptions=MyException)
+    def _():
+        nonlocal tried
+        tried += 1
+        if tried == 2:
+            assert True
+        else:
+            raise MyException
+
+
+def test_call_custom_exception_with_tuple():
+    tried = 0
+
+    @eventually_consistent.call(
+        tries=3, exceptions=(MyException, AssertionError))
+    def _():
+        nonlocal tried
+        tried += 1
+        if tried == 3:
+            assert True
+        elif tried == 2:
+            assert False
+        else:
+            raise MyException