diff options
Diffstat (limited to 'test/test_client_async.py')
| -rw-r--r-- | test/test_client_async.py | 44 |
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() |