From 4b7330ddf7441f839be02dc9dd3eecb3ff7909a3 Mon Sep 17 00:00:00 2001
From: Jim Witschey <jim.witschey@gmail.com>
Date: Fri, 21 Apr 2017 17:16:06 -0400
Subject: [PATCH] fixes for ExponentialReconnectionPolicy

- documentation specified floats for base_delay and max_delay; fix tests
  to match
- fix issue from PYTHON-707
---
 cassandra/policies.py       | 12 ++++++++++--
 tests/unit/test_policies.py | 18 ++++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index 9be0125f..108dcbaa 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -564,9 +564,17 @@ class ExponentialReconnectionPolicy(ReconnectionPolicy):
         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
index 494b2e05..f1f4995d 100644
--- a/tests/unit/test_policies.py
+++ b/tests/unit/test_policies.py
@@ -885,8 +885,8 @@ class ExponentialReconnectionPolicyTest(unittest.TestCase):
         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 @@ class ExponentialReconnectionPolicyTest(unittest.TestCase):
         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 @@ class ExponentialReconnectionPolicyTest(unittest.TestCase):
             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