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

Author: dpkp

kafka/conn.py

@@ -594,9 +594,16 @@ def blacked_out(self):
         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:

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