X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/a942e37250873d383bd885ba0dba70c63b3c073d..b20590222beddb52c8c89294ed3a324c8c7190a2:/sdk/python/tests/test_retry.py

diff --git a/sdk/python/tests/test_retry.py b/sdk/python/tests/test_retry.py
index ed0a406e85..0c1110c5ce 100644
--- a/sdk/python/tests/test_retry.py
+++ b/sdk/python/tests/test_retry.py
@@ -7,7 +7,7 @@ import arvados.errors as arv_error
 import arvados.retry as arv_retry
 import mock
 
-from arvados_testutil import fake_httplib2_response
+from arvados_testutil import fake_requests_response
 
 class RetryLoopTestMixin(object):
     @staticmethod
@@ -150,7 +150,7 @@ class RetryLoopBackoffTestCase(unittest.TestCase, RetryLoopTestMixin):
 class CheckHTTPResponseSuccessTestCase(unittest.TestCase):
     def results_map(self, *codes):
         for code in codes:
-            response = (fake_httplib2_response(code), None)
+            response = fake_requests_response(code, None)
             yield code, arv_retry.check_http_response_success(response)
 
     def check(assert_name):
@@ -194,5 +194,32 @@ class CheckHTTPResponseSuccessTestCase(unittest.TestCase):
         self.check_is(None, 0, 99, 600, -200)
 
 
+class RetryMethodTestCase(unittest.TestCase):
+    class Tester(object):
+        def __init__(self):
+            self.num_retries = 1
+
+        @arv_retry.retry_method
+        def check(self, a, num_retries=None, z=0):
+            return (a, num_retries, z)
+
+
+    def test_positional_arg_passed(self):
+        self.assertEqual((3, 2, 0), self.Tester().check(3, 2))
+
+    def test_keyword_arg_passed(self):
+        self.assertEqual((4, 3, 0), self.Tester().check(num_retries=3, a=4))
+
+    def test_not_specified(self):
+        self.assertEqual((0, 1, 0), self.Tester().check(0))
+
+    def test_not_specified_with_other_kwargs(self):
+        self.assertEqual((1, 1, 1), self.Tester().check(1, z=1))
+
+    def test_bad_call(self):
+        with self.assertRaises(TypeError):
+            self.Tester().check(num_retries=2)
+
+
 if __name__ == '__main__':
     unittest.main()