Fix BrokerConnection.connection_delay() to return milliseconds (#1414)

2 files changed, 18 insertions, 2 deletions

diff --git a/kafka/conn.py b/kafka/conn.py
index 0c8ae9a..d778c31 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -594,9 +594,16 @@ class BrokerConnection(object):
         return False
 
     def connection_delay(self):
-        time_waited_ms = time.time() - (self.last_attempt or 0)
+        """
+        Return the number of milliseconds to wait, based on the connection
+        state, before attempting to send data. When disconnected, this respects
+        the reconnect backoff time. When connecting, returns 0 to allow
+        non-blocking connect to finish. When connected, returns a very large
+        number to handle slow/stalled connections.
+        """
+        time_waited = time.time() - (self.last_attempt or 0)
         if self.state is ConnectionStates.DISCONNECTED:
-            return max(self._reconnect_backoff - time_waited_ms, 0)
+            return max(self._reconnect_backoff - time_waited, 0) * 1000
         elif self.connecting():
             return 0
         else:
diff --git a/test/test_conn.py b/test/test_conn.py
index 5698519..f35cebe 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -72,6 +72,15 @@ def test_blacked_out(conn):
     assert conn.blacked_out() is True
 
 
+def test_connection_delay(conn):
+    conn.last_attempt = time.time()
+    assert round(conn.connection_delay()) == round(conn.config['reconnect_backoff_ms'])
+    conn.state = ConnectionStates.CONNECTING
+    assert conn.connection_delay() == 0
+    conn.state = ConnectionStates.CONNECTED
+    assert conn.connection_delay() == float('inf')
+
+
 def test_connected(conn):
     assert conn.connected() is False
     conn.state = ConnectionStates.CONNECTED