Merge pull request #639 from dpkp/conn_state_callback

Use KafkaClient callback to manage BrokerConnection state changes

1 files changed, 32 insertions, 12 deletions

diff --git a/test/test_client_async.py b/test/test_client_async.py
index 6da5394..ad76aad 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -1,5 +1,5 @@
-import time
 import socket
+import time
 
 import pytest
 
@@ -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
@@ -83,26 +89,40 @@ def test_maybe_connect(conn):
     else:
         assert False, 'Exception not raised'
 
+    # New node_id creates a conn object
     assert 0 not in cli._conns
     conn.state = ConnectionStates.DISCONNECTED
     conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.CONNECTING)
     assert cli._maybe_connect(0) is False
     assert cli._conns[0] is conn
-    assert 0 in cli._connecting
 
-    conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.CONNECTED)
-    assert cli._maybe_connect(0) is True
-    assert 0 not in cli._connecting
+
+def test_conn_state_change(mocker, conn):
+    cli = KafkaClient()
+
+    node_id = 0
+    conn.state = ConnectionStates.CONNECTING
+    cli._conn_state_change(node_id, conn)
+    assert node_id in cli._connecting
+
+    conn.state = ConnectionStates.CONNECTED
+    cli._conn_state_change(node_id, conn)
+    assert node_id not in cli._connecting
 
     # Failure to connect should trigger metadata update
     assert cli.cluster._need_update is False
-    conn.state = ConnectionStates.CONNECTING
-    cli._connecting.add(0)
-    conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.DISCONNECTED)
-    assert cli._maybe_connect(0) is False
-    assert 0 not in cli._connecting
+    conn.state = ConnectionStates.DISCONNECTING
+    cli._conn_state_change(node_id, conn)
+    assert node_id not in cli._connecting
     assert cli.cluster._need_update is True
 
+    conn.state = ConnectionStates.CONNECTING
+    cli._conn_state_change(node_id, conn)
+    assert node_id in cli._connecting
+    conn.state = ConnectionStates.DISCONNECTING
+    cli._conn_state_change(node_id, conn)
+    assert node_id not in cli._connecting
+
 
 def test_ready(mocker, conn):
     cli = KafkaClient()