diff options
| -rw-r--r-- | kafka/client_async.py | 5 | ||||
| -rw-r--r-- | test/test_client_async.py | 10 |
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 |