Add state_change_callback to bootstrap connectionconn_state_callback

2 files changed, 12 insertions, 3 deletions

diff --git a/kafka/client_async.py b/kafka/client_async.py
index d0a3723..0c22f90 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -119,7 +119,10 @@ class KafkaClient(object):
         metadata_request = MetadataRequest[0]([])
         for host, port, afi in hosts:
             log.debug("Attempting to bootstrap via node at %s:%s", host, port)
-            bootstrap = BrokerConnection(host, port, afi, **self.config)
+            cb = functools.partial(self._conn_state_change, 'bootstrap')
+            bootstrap = BrokerConnection(host, port, afi,
+                                         state_change_callback=cb,
+                                         **self.config)
             bootstrap.connect()
             while bootstrap.connecting():
                 bootstrap.connect()
diff --git a/test/test_client_async.py b/test/test_client_async.py
index ae8549d..ad76aad 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -34,7 +34,10 @@ def test_bootstrap_servers(mocker, bootstrap, expected_hosts):
 def test_bootstrap_success(conn):
     conn.state = ConnectionStates.CONNECTED
     cli = KafkaClient()
-    conn.assert_called_once_with('localhost', 9092, socket.AF_INET, **cli.config)
+    args, kwargs = conn.call_args
+    assert args == ('localhost', 9092, socket.AF_INET)
+    kwargs.pop('state_change_callback')
+    assert kwargs == cli.config
     conn.connect.assert_called_with()
     conn.send.assert_called_once_with(MetadataRequest[0]([]))
     assert cli._bootstrap_fails == 0
@@ -44,7 +47,10 @@ def test_bootstrap_success(conn):
 def test_bootstrap_failure(conn):
     conn.state = ConnectionStates.DISCONNECTED
     cli = KafkaClient()
-    conn.assert_called_once_with('localhost', 9092, socket.AF_INET, **cli.config)
+    args, kwargs = conn.call_args
+    assert args == ('localhost', 9092, socket.AF_INET)
+    kwargs.pop('state_change_callback')
+    assert kwargs == cli.config
     conn.connect.assert_called_with()
     conn.close.assert_called_with()
     assert cli._bootstrap_fails == 1