Add BrokerConnection.connect_blocking() (#1411)

2 files changed, 6 insertions, 3 deletions

diff --git a/test/conftest.py b/test/conftest.py
index d53ff23..52ebfb4 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -128,6 +128,7 @@ def conn(mocker):
         return state
     conn._set_conn_state = _set_conn_state
     conn.connect.side_effect = lambda: conn.state
+    conn.connect_blocking.return_value = True
     conn.connecting = lambda: conn.state in (ConnectionStates.CONNECTING,
                                              ConnectionStates.HANDSHAKE)
     conn.connected = lambda: conn.state is ConnectionStates.CONNECTED
diff --git a/test/test_client_async.py b/test/test_client_async.py
index eece139..eccb564 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -55,21 +55,22 @@ def test_bootstrap_success(conn):
     kwargs.pop('state_change_callback')
     kwargs.pop('node_id')
     assert kwargs == cli.config
-    conn.connect.assert_called_with()
+    conn.connect_blocking.assert_called_with()
     conn.send.assert_called_once_with(MetadataRequest[0]([]))
     assert cli._bootstrap_fails == 0
     assert cli.cluster.brokers() == set([BrokerMetadata(0, 'foo', 12, None),
                                          BrokerMetadata(1, 'bar', 34, None)])
 
+
 def test_bootstrap_failure(conn):
-    conn.state = ConnectionStates.DISCONNECTED
+    conn.connect_blocking.return_value = False
     cli = KafkaClient(api_version=(0, 9))
     args, kwargs = conn.call_args
     assert args == ('localhost', 9092, socket.AF_UNSPEC)
     kwargs.pop('state_change_callback')
     kwargs.pop('node_id')
     assert kwargs == cli.config
-    conn.connect.assert_called_with()
+    conn.connect_blocking.assert_called_with()
     conn.close.assert_called_with()
     assert cli._bootstrap_fails == 1
     assert cli.cluster.brokers() == set()
@@ -95,6 +96,7 @@ def test_can_connect(cli, conn):
     conn.blacked_out.return_value = True
     assert not cli._can_connect(0)
 
+
 def test_maybe_connect(cli, conn):
     try:
         # Node not in metadata, raises AssertionError