fixes for ExponentialReconnectionPolicy

mambocab · mambocab · commit 4b7330ddf744 · 2017-04-24T12:39:03.000-04:00
- documentation specified floats for base_delay and max_delay; fix tests
  to match
- fix issue from PYTHON-707
diff --git a/cassandra/policies.py b/cassandra/policies.py
@@ -564,9 +564,17 @@ def __init__(self, base_delay, max_delay, max_attempts=64):
         self.max_attempts = max_attempts
 
     def new_schedule(self):
-        i = 0
+        i, overflowed = 0, False
         while self.max_attempts is None or i < self.max_attempts:
-            yield min(self.base_delay * (2 ** i), self.max_delay)
+            if overflowed:
+                yield self.max_delay
+            else:
+                try:
+                    yield min(self.base_delay * (2 ** i), self.max_delay)
+                except OverflowError:
+                    overflowed = True
+                    yield self.max_delay
+
             i += 1
 
 
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
@@ -885,8 +885,8 @@ def test_bad_vals(self):
         self.assertRaises(ValueError, ExponentialReconnectionPolicy, 1, 2,-1)
 
     def test_schedule_no_max(self):
-        base_delay = 2
-        max_delay = 100
+        base_delay = 2.0
+        max_delay = 100.0
         test_iter = 10000
         policy = ExponentialReconnectionPolicy(base_delay=base_delay, max_delay=max_delay, max_attempts=None)
         sched_slice = list(islice(policy.new_schedule(), 0, test_iter))
@@ -895,8 +895,8 @@ def test_schedule_no_max(self):
         self.assertEqual(len(sched_slice), test_iter)
 
     def test_schedule_with_max(self):
-        base_delay = 2
-        max_delay = 100
+        base_delay = 2.0
+        max_delay = 100.0
         max_attempts = 64
         policy = ExponentialReconnectionPolicy(base_delay=base_delay, max_delay=max_delay, max_attempts=max_attempts)
         schedule = list(policy.new_schedule())
@@ -909,6 +909,16 @@ def test_schedule_with_max(self):
             else:
                 self.assertEqual(delay, max_delay)
 
+    def test_schedule_exactly_one_attempt(self):
+        base_delay = 2.0
+        max_delay = 100.0
+        max_attempts = 1
+        policy = ExponentialReconnectionPolicy(
+            base_delay=base_delay, max_delay=max_delay, max_attempts=max_attempts
+        )
+        self.assertEqual(len(list(policy.new_schedule())), 1)
+
+
 ONE = ConsistencyLevel.ONE
 
 
