From 95d244ec7c0ac3def97b4de21aa3ab609fa2b437 Mon Sep 17 00:00:00 2001
From: Charles-Henri de Boysson <ceache@users.noreply.github.com>
Date: Thu, 29 Feb 2024 01:44:58 -0500
Subject: [PATCH] fix(core): Proper retry count in KazooRetry

Make sure the number of attempts matches the `max_retry` parameter.
Add unit tests to that effect.
---
 kazoo/retry.py            |   2 +-
 kazoo/tests/test_retry.py | 140 +++++++++++++++++++++-----------------
 2 files changed, 80 insertions(+), 62 deletions(-)

diff --git a/kazoo/retry.py b/kazoo/retry.py
index a51a1b56..fb9e8fc7 100644
--- a/kazoo/retry.py
+++ b/kazoo/retry.py
@@ -133,10 +133,10 @@ def __call__(self, func, *args, **kwargs):
             except ConnectionClosedError:
                 raise
             except self.retry_exceptions:
+                self._attempts += 1
                 # Note: max_tries == -1 means infinite tries.
                 if self._attempts == self.max_tries:
                     raise RetryFailedError("Too many retry attempts")
-                self._attempts += 1
                 jitter = random.uniform(
                     1.0 - self.max_jitter, 1.0 + self.max_jitter
                 )
diff --git a/kazoo/tests/test_retry.py b/kazoo/tests/test_retry.py
index 5b79eca9..e6b360b3 100644
--- a/kazoo/tests/test_retry.py
+++ b/kazoo/tests/test_retry.py
@@ -1,89 +1,107 @@
-import unittest
-
 import pytest
 
 
-class TestRetrySleeper(unittest.TestCase):
-    def _pass(self):
-        pass
+def _sleep_func(_time):
+    pass
+
 
-    def _fail(self, times=1):
-        from kazoo.retry import ForceRetryError
+def _makeOne(*args, **kwargs):
+    from kazoo.retry import KazooRetry
 
-        scope = dict(times=0)
+    return KazooRetry(*args, sleep_func=_sleep_func, **kwargs)
 
-        def inner():
-            if scope["times"] >= times:
-                pass
-            else:
-                scope["times"] += 1
-                raise ForceRetryError("Failed!")
 
-        return inner
+_try_counts = 0
 
-    def _makeOne(self, *args, **kwargs):
-        from kazoo.retry import KazooRetry
 
-        return KazooRetry(*args, **kwargs)
+def _make_try_fun(times=1):
+    """Returns a function that raises ForceRetryError `times` time before
+    returning None."""
+    from kazoo.retry import ForceRetryError
 
-    def test_reset(self):
-        retry = self._makeOne(delay=0, max_tries=2)
-        retry(self._fail())
-        assert retry._attempts == 1
-        retry.reset()
-        assert retry._attempts == 0
+    global _try_counts
 
-    def test_too_many_tries(self):
-        from kazoo.retry import RetryFailedError
+    _try_counts = 0
 
-        retry = self._makeOne(delay=0)
-        with pytest.raises(RetryFailedError):
-            retry(self._fail(times=999))
-        assert retry._attempts == 1
+    def inner():
+        global _try_counts
 
-    def test_maximum_delay(self):
-        def sleep_func(_time):
+        if _try_counts >= times:
             pass
+        else:
+            _try_counts += 1
+            raise ForceRetryError("Failed!")
+
+    return inner
+
+
+def test_reset():
+    global _try_counts
+
+    retry = _makeOne(delay=0, max_tries=2)
+    retry(_make_try_fun())
+    assert _try_counts == retry._attempts == 1
+    retry.reset()
+    assert retry._attempts == 0
+
+
+def test_too_many_tries():
+    global _try_counts
+
+    from kazoo.retry import RetryFailedError
+
+    retry = _makeOne(delay=0, max_tries=10)
+    with pytest.raises(RetryFailedError):
+        retry(_make_try_fun(times=999))
+    assert _try_counts == retry._attempts == 10
+
+
+def test_maximum_delay():
+    global _try_counts
+
+    retry = _makeOne(delay=10, max_tries=100, max_jitter=0)
+    retry(_make_try_fun(times=2))
+    assert _try_counts == 2
+    assert retry._cur_delay == 10 * 2**2, "Normal exponential backoff"
+    retry.reset()
+    retry(_make_try_fun(times=10))
+    assert _try_counts == 10
+    assert retry._cur_delay == 60, "Delay capped by maximun"
+    # gevent's sleep function is picky about the type
+    assert isinstance(retry._cur_delay, float)
 
-        retry = self._makeOne(delay=10, max_tries=100, sleep_func=sleep_func)
-        retry(self._fail(times=10))
-        assert retry._cur_delay < 4000
-        # gevent's sleep function is picky about the type
-        assert type(retry._cur_delay) == float
 
-    def test_copy(self):
-        def _sleep(t):
-            return None
+def test_copy():
+    retry = _makeOne()
+    rcopy = retry.copy()
+    assert rcopy.sleep_func is retry.sleep_func
 
-        retry = self._makeOne(sleep_func=_sleep)
-        rcopy = retry.copy()
-        assert rcopy.sleep_func is _sleep
 
+def test_connection_closed():
+    from kazoo.exceptions import ConnectionClosedError
 
-class TestKazooRetry(unittest.TestCase):
-    def _makeOne(self, **kw):
-        from kazoo.retry import KazooRetry
+    retry = _makeOne()
 
-        return KazooRetry(**kw)
+    def testit():
+        raise ConnectionClosedError()
 
-    def test_connection_closed(self):
-        from kazoo.exceptions import ConnectionClosedError
+    with pytest.raises(ConnectionClosedError):
+        retry(testit)
 
-        retry = self._makeOne()
 
-        def testit():
-            raise ConnectionClosedError()
+def test_session_expired():
+    from kazoo.exceptions import SessionExpiredError
+    from kazoo.retry import RetryFailedError
 
-        with pytest.raises(ConnectionClosedError):
-            retry(testit)
+    retry = _makeOne(max_tries=1)
 
-    def test_session_expired(self):
-        from kazoo.exceptions import SessionExpiredError
+    def testit():
+        raise SessionExpiredError()
 
-        retry = self._makeOne(max_tries=1)
+    with pytest.raises(RetryFailedError) as err:
+        retry(testit)
 
-        def testit():
-            raise SessionExpiredError()
+    retry = _makeOne(max_tries=1, ignore_expire=False)
 
-        with pytest.raises(Exception):
-            retry(testit)
+    with pytest.raises(SessionExpiredError) as err:
+        retry(testit)