diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/test_client_async.py | 84 | ||||
| -rw-r--r-- | test/test_conn.py | 318 | ||||
| -rw-r--r-- | test/test_conn_legacy.py | 242 |
3 files changed, 361 insertions, 283 deletions
diff --git a/test/test_client_async.py b/test/test_client_async.py index e0b98c4..884686d 100644 --- a/test/test_client_async.py +++ b/test/test_client_async.py @@ -41,7 +41,8 @@ def conn(mocker): [(0, 'foo', 12), (1, 'bar', 34)], # brokers [])) # topics conn.blacked_out.return_value = False - conn.connect.return_value = conn.state + conn.connect.side_effect = lambda: conn.state + conn.connected = lambda: conn.connect() is ConnectionStates.CONNECTED return conn @@ -76,7 +77,7 @@ def test_can_connect(conn): assert cli._can_connect(0) # Node is connected, can't reconnect - cli._initiate_connect(0) + assert cli._maybe_connect(0) is True assert not cli._can_connect(0) # Node is disconnected, can connect @@ -87,60 +88,47 @@ def test_can_connect(conn): conn.blacked_out.return_value = True assert not cli._can_connect(0) -def test_initiate_connect(conn): +def test_maybe_connect(conn): cli = KafkaClient() try: # Node not in metadata, raises AssertionError - cli._initiate_connect(2) + cli._maybe_connect(2) except AssertionError: pass else: assert False, 'Exception not raised' assert 0 not in cli._conns - state = cli._initiate_connect(0) + conn.state = ConnectionStates.DISCONNECTED + conn.connect.side_effect = lambda: ConnectionStates.CONNECTING + assert cli._maybe_connect(0) is False assert cli._conns[0] is conn - assert state is conn.state - - -def test_finish_connect(conn): - cli = KafkaClient() - try: - # Node not in metadata, raises AssertionError - cli._initiate_connect(2) - except AssertionError: - pass - else: - assert False, 'Exception not raised' - - assert 0 not in cli._conns - cli._initiate_connect(0) - - conn.connect.return_value = ConnectionStates.CONNECTING - state = cli._finish_connect(0) assert 0 in cli._connecting - assert state is ConnectionStates.CONNECTING - conn.connect.return_value = ConnectionStates.CONNECTED - state = cli._finish_connect(0) + conn.state = ConnectionStates.CONNECTING + conn.connect.side_effect = lambda: ConnectionStates.CONNECTED + assert cli._maybe_connect(0) is True assert 0 not in cli._connecting - assert state is ConnectionStates.CONNECTED # Failure to connect should trigger metadata update - assert not cli.cluster._need_update + assert cli.cluster._need_update is False cli._connecting.add(0) - conn.connect.return_value = ConnectionStates.DISCONNECTED - state = cli._finish_connect(0) + conn.state = ConnectionStates.CONNECTING + conn.connect.side_effect = lambda: ConnectionStates.DISCONNECTED + assert cli._maybe_connect(0) is False assert 0 not in cli._connecting - assert state is ConnectionStates.DISCONNECTED - assert cli.cluster._need_update + assert cli.cluster._need_update is True def test_ready(conn): cli = KafkaClient() - # Node not in metadata - assert not cli.ready(2) + # Node not in metadata raises Exception + try: + cli.ready(2) + assert False, 'Exception not raised' + except AssertionError: + pass # Node in metadata will connect assert 0 not in cli._conns @@ -176,13 +164,13 @@ def test_ready(conn): # disconnected nodes, not ready assert cli.ready(0) assert cli.ready(1) - conn.connected.return_value = False + conn.state = ConnectionStates.DISCONNECTED assert not cli.ready(0) - conn.connected.return_value = True # connecting node connects cli._connecting.add(0) - conn.connected.return_value = False + conn.state = ConnectionStates.CONNECTING + conn.connect.side_effect = lambda: ConnectionStates.CONNECTED cli.ready(0) assert 0 not in cli._connecting assert cli._conns[0].connect.called_with() @@ -195,13 +183,13 @@ def test_close(conn): cli.close(2) # Single node close - cli._initiate_connect(0) + cli._maybe_connect(0) assert not conn.close.call_count cli.close(0) assert conn.close.call_count == 1 # All node close - cli._initiate_connect(1) + cli._maybe_connect(1) cli.close() assert conn.close.call_count == 3 @@ -213,7 +201,7 @@ def test_is_disconnected(conn): conn.state = ConnectionStates.DISCONNECTED assert not cli.is_disconnected(0) - cli._initiate_connect(0) + cli._maybe_connect(0) assert cli.is_disconnected(0) conn.state = ConnectionStates.CONNECTING @@ -225,14 +213,22 @@ def test_is_disconnected(conn): def test_send(conn): cli = KafkaClient() + + # Send to unknown node => raises AssertionError try: cli.send(2, None) - except Errors.NodeNotReadyError: + assert False, 'Exception not raised' + except AssertionError: pass - else: - assert False, 'NodeNotReadyError not raised' - cli._initiate_connect(0) + # Send to disconnected node => NodeNotReady + conn.state = ConnectionStates.DISCONNECTED + f = cli.send(0, None) + assert f.failed() + assert isinstance(f.exception, Errors.NodeNotReadyError) + + conn.state = ConnectionStates.CONNECTED + cli._maybe_connect(0) # ProduceRequest w/ 0 required_acks -> no response request = ProduceRequest(0, 0, []) ret = cli.send(0, request) diff --git a/test/test_conn.py b/test/test_conn.py index f0ef8fb..d394f74 100644 --- a/test/test_conn.py +++ b/test/test_conn.py @@ -1,242 +1,82 @@ +# pylint: skip-file +from __future__ import absolute_import + +from errno import EALREADY, EINPROGRESS, EISCONN, ECONNRESET import socket -import struct -from threading import Thread - -import mock -from . import unittest - -from kafka.common import ConnectionError -from kafka.conn import KafkaConnection, collect_hosts, DEFAULT_SOCKET_TIMEOUT_SECONDS - -class ConnTest(unittest.TestCase): - def setUp(self): - - self.config = { - 'host': 'localhost', - 'port': 9090, - 'request_id': 0, - 'payload': b'test data', - 'payload2': b'another packet' - } - - # Mocking socket.create_connection will cause _sock to always be a - # MagicMock() - patcher = mock.patch('socket.create_connection', spec=True) - self.MockCreateConn = patcher.start() - self.addCleanup(patcher.stop) - - # Also mock socket.sendall() to appear successful - self.MockCreateConn().sendall.return_value = None - - # And mock socket.recv() to return two payloads, then '', then raise - # Note that this currently ignores the num_bytes parameter to sock.recv() - payload_size = len(self.config['payload']) - payload2_size = len(self.config['payload2']) - self.MockCreateConn().recv.side_effect = [ - struct.pack('>i', payload_size), - struct.pack('>%ds' % payload_size, self.config['payload']), - struct.pack('>i', payload2_size), - struct.pack('>%ds' % payload2_size, self.config['payload2']), - b'' - ] - - # Create a connection object - self.conn = KafkaConnection(self.config['host'], self.config['port']) - - # Reset any mock counts caused by __init__ - self.MockCreateConn.reset_mock() - - def test_collect_hosts__happy_path(self): - hosts = "localhost:1234,localhost" - results = collect_hosts(hosts) - - self.assertEqual(set(results), set([ - ('localhost', 1234, socket.AF_INET), - ('localhost', 9092, socket.AF_INET), - ])) - - def test_collect_hosts__ipv6(self): - hosts = "[localhost]:1234,[2001:1000:2000::1],[2001:1000:2000::1]:1234" - results = collect_hosts(hosts) - - self.assertEqual(set(results), set([ - ('localhost', 1234, socket.AF_INET6), - ('2001:1000:2000::1', 9092, socket.AF_INET6), - ('2001:1000:2000::1', 1234, socket.AF_INET6), - ])) - - def test_collect_hosts__string_list(self): - hosts = [ - 'localhost:1234', - 'localhost', - '[localhost]', - '2001::1', - '[2001::1]:1234', - ] - - results = collect_hosts(hosts) - - self.assertEqual(set(results), set([ - ('localhost', 1234, socket.AF_INET), - ('localhost', 9092, socket.AF_INET), - ('localhost', 9092, socket.AF_INET6), - ('2001::1', 9092, socket.AF_INET6), - ('2001::1', 1234, socket.AF_INET6), - ])) - - def test_collect_hosts__with_spaces(self): - hosts = "localhost:1234, localhost" - results = collect_hosts(hosts) - - self.assertEqual(set(results), set([ - ('localhost', 1234, socket.AF_INET), - ('localhost', 9092, socket.AF_INET), - ])) - - - def test_send(self): - self.conn.send(self.config['request_id'], self.config['payload']) - self.conn._sock.sendall.assert_called_with(self.config['payload']) - - def test_init_creates_socket_connection(self): - KafkaConnection(self.config['host'], self.config['port']) - self.MockCreateConn.assert_called_with((self.config['host'], self.config['port']), DEFAULT_SOCKET_TIMEOUT_SECONDS) - - def test_init_failure_raises_connection_error(self): - - def raise_error(*args): - raise socket.error - - assert socket.create_connection is self.MockCreateConn - socket.create_connection.side_effect=raise_error - with self.assertRaises(ConnectionError): - KafkaConnection(self.config['host'], self.config['port']) - - def test_send__reconnects_on_dirty_conn(self): - - # Dirty the connection - try: - self.conn._raise_connection_error() - except ConnectionError: - pass - - # Now test that sending attempts to reconnect - self.assertEqual(self.MockCreateConn.call_count, 0) - self.conn.send(self.config['request_id'], self.config['payload']) - self.assertEqual(self.MockCreateConn.call_count, 1) - - def test_send__failure_sets_dirty_connection(self): - - def raise_error(*args): - raise socket.error - - assert isinstance(self.conn._sock, mock.Mock) - self.conn._sock.sendall.side_effect=raise_error - try: - self.conn.send(self.config['request_id'], self.config['payload']) - except ConnectionError: - self.assertIsNone(self.conn._sock) - - def test_recv(self): - - self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload']) - - def test_recv__reconnects_on_dirty_conn(self): - - # Dirty the connection - try: - self.conn._raise_connection_error() - except ConnectionError: - pass - - # Now test that recv'ing attempts to reconnect - self.assertEqual(self.MockCreateConn.call_count, 0) - self.conn.recv(self.config['request_id']) - self.assertEqual(self.MockCreateConn.call_count, 1) - - def test_recv__failure_sets_dirty_connection(self): - - def raise_error(*args): - raise socket.error - - # test that recv'ing attempts to reconnect - assert isinstance(self.conn._sock, mock.Mock) - self.conn._sock.recv.side_effect=raise_error - try: - self.conn.recv(self.config['request_id']) - except ConnectionError: - self.assertIsNone(self.conn._sock) - - def test_recv__doesnt_consume_extra_data_in_stream(self): - - # Here just test that each call to recv will return a single payload - self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload']) - self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload2']) - - def test_get_connected_socket(self): - s = self.conn.get_connected_socket() - - self.assertEqual(s, self.MockCreateConn()) - - def test_get_connected_socket_on_dirty_conn(self): - # Dirty the connection - try: - self.conn._raise_connection_error() - except ConnectionError: - pass - - # Test that get_connected_socket tries to connect - self.assertEqual(self.MockCreateConn.call_count, 0) - self.conn.get_connected_socket() - self.assertEqual(self.MockCreateConn.call_count, 1) +import time + +import pytest + +from kafka.conn import BrokerConnection, ConnectionStates + + +@pytest.fixture +def socket(mocker): + socket = mocker.MagicMock() + socket.connect_ex.return_value = 0 + mocker.patch('socket.socket', return_value=socket) + return socket + + +@pytest.fixture +def conn(socket): + conn = BrokerConnection('localhost', 9092, socket.AF_INET) + return conn + + +@pytest.mark.parametrize("states", [ + (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING),), + (([EALREADY, EALREADY], ConnectionStates.CONNECTING),), + (([0], ConnectionStates.CONNECTED),), + (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING), + ([ECONNRESET], ConnectionStates.DISCONNECTED)), + (([EINPROGRESS, EALREADY], ConnectionStates.CONNECTING), + ([EALREADY], ConnectionStates.CONNECTING), + ([EISCONN], ConnectionStates.CONNECTED)), +]) +def test_connect(socket, conn, states): + assert conn.state is ConnectionStates.DISCONNECTED + + for errno, state in states: + socket.connect_ex.side_effect = errno + conn.connect() + assert conn.state is state + + +def test_connect_timeout(socket, conn): + assert conn.state is ConnectionStates.DISCONNECTED + + # Initial connect returns EINPROGRESS + # immediate inline connect returns EALREADY + # second explicit connect returns EALREADY + # third explicit connect returns EALREADY and times out via last_attempt + socket.connect_ex.side_effect = [EINPROGRESS, EALREADY, EALREADY, EALREADY] + conn.connect() + assert conn.state is ConnectionStates.CONNECTING + conn.connect() + assert conn.state is ConnectionStates.CONNECTING + conn.last_attempt = 0 + conn.connect() + assert conn.state is ConnectionStates.DISCONNECTED + + +def test_blacked_out(conn): + assert not conn.blacked_out() + conn.last_attempt = time.time() + assert conn.blacked_out() + + +def test_connected(conn): + assert not conn.connected() + conn.state = ConnectionStates.CONNECTED + assert conn.connected() + - def test_close__object_is_reusable(self): +def test_connecting(conn): + assert not conn.connecting() + conn.state = ConnectionStates.CONNECTING + assert conn.connecting() + conn.state = ConnectionStates.CONNECTED + assert not conn.connecting() - # test that sending to a closed connection - # will re-connect and send data to the socket - self.conn.close() - self.conn.send(self.config['request_id'], self.config['payload']) - self.assertEqual(self.MockCreateConn.call_count, 1) - self.conn._sock.sendall.assert_called_with(self.config['payload']) - - -class TestKafkaConnection(unittest.TestCase): - @mock.patch('socket.create_connection') - def test_copy(self, socket): - """KafkaConnection copies work as expected""" - - conn = KafkaConnection('kafka', 9092) - self.assertEqual(socket.call_count, 1) - - copy = conn.copy() - self.assertEqual(socket.call_count, 1) - self.assertEqual(copy.host, 'kafka') - self.assertEqual(copy.port, 9092) - self.assertEqual(copy._sock, None) - - copy.reinit() - self.assertEqual(socket.call_count, 2) - self.assertNotEqual(copy._sock, None) - - @mock.patch('socket.create_connection') - def test_copy_thread(self, socket): - """KafkaConnection copies work in other threads""" - - err = [] - copy = KafkaConnection('kafka', 9092).copy() - - def thread_func(err, copy): - try: - self.assertEqual(copy.host, 'kafka') - self.assertEqual(copy.port, 9092) - self.assertNotEqual(copy._sock, None) - except Exception as e: - err.append(e) - else: - err.append(None) - thread = Thread(target=thread_func, args=(err, copy)) - thread.start() - thread.join() - - self.assertEqual(err, [None]) - self.assertEqual(socket.call_count, 2) +# TODO: test_send, test_recv, test_can_send_more, test_close diff --git a/test/test_conn_legacy.py b/test/test_conn_legacy.py new file mode 100644 index 0000000..f0ef8fb --- /dev/null +++ b/test/test_conn_legacy.py @@ -0,0 +1,242 @@ +import socket +import struct +from threading import Thread + +import mock +from . import unittest + +from kafka.common import ConnectionError +from kafka.conn import KafkaConnection, collect_hosts, DEFAULT_SOCKET_TIMEOUT_SECONDS + +class ConnTest(unittest.TestCase): + def setUp(self): + + self.config = { + 'host': 'localhost', + 'port': 9090, + 'request_id': 0, + 'payload': b'test data', + 'payload2': b'another packet' + } + + # Mocking socket.create_connection will cause _sock to always be a + # MagicMock() + patcher = mock.patch('socket.create_connection', spec=True) + self.MockCreateConn = patcher.start() + self.addCleanup(patcher.stop) + + # Also mock socket.sendall() to appear successful + self.MockCreateConn().sendall.return_value = None + + # And mock socket.recv() to return two payloads, then '', then raise + # Note that this currently ignores the num_bytes parameter to sock.recv() + payload_size = len(self.config['payload']) + payload2_size = len(self.config['payload2']) + self.MockCreateConn().recv.side_effect = [ + struct.pack('>i', payload_size), + struct.pack('>%ds' % payload_size, self.config['payload']), + struct.pack('>i', payload2_size), + struct.pack('>%ds' % payload2_size, self.config['payload2']), + b'' + ] + + # Create a connection object + self.conn = KafkaConnection(self.config['host'], self.config['port']) + + # Reset any mock counts caused by __init__ + self.MockCreateConn.reset_mock() + + def test_collect_hosts__happy_path(self): + hosts = "localhost:1234,localhost" + results = collect_hosts(hosts) + + self.assertEqual(set(results), set([ + ('localhost', 1234, socket.AF_INET), + ('localhost', 9092, socket.AF_INET), + ])) + + def test_collect_hosts__ipv6(self): + hosts = "[localhost]:1234,[2001:1000:2000::1],[2001:1000:2000::1]:1234" + results = collect_hosts(hosts) + + self.assertEqual(set(results), set([ + ('localhost', 1234, socket.AF_INET6), + ('2001:1000:2000::1', 9092, socket.AF_INET6), + ('2001:1000:2000::1', 1234, socket.AF_INET6), + ])) + + def test_collect_hosts__string_list(self): + hosts = [ + 'localhost:1234', + 'localhost', + '[localhost]', + '2001::1', + '[2001::1]:1234', + ] + + results = collect_hosts(hosts) + + self.assertEqual(set(results), set([ + ('localhost', 1234, socket.AF_INET), + ('localhost', 9092, socket.AF_INET), + ('localhost', 9092, socket.AF_INET6), + ('2001::1', 9092, socket.AF_INET6), + ('2001::1', 1234, socket.AF_INET6), + ])) + + def test_collect_hosts__with_spaces(self): + hosts = "localhost:1234, localhost" + results = collect_hosts(hosts) + + self.assertEqual(set(results), set([ + ('localhost', 1234, socket.AF_INET), + ('localhost', 9092, socket.AF_INET), + ])) + + + def test_send(self): + self.conn.send(self.config['request_id'], self.config['payload']) + self.conn._sock.sendall.assert_called_with(self.config['payload']) + + def test_init_creates_socket_connection(self): + KafkaConnection(self.config['host'], self.config['port']) + self.MockCreateConn.assert_called_with((self.config['host'], self.config['port']), DEFAULT_SOCKET_TIMEOUT_SECONDS) + + def test_init_failure_raises_connection_error(self): + + def raise_error(*args): + raise socket.error + + assert socket.create_connection is self.MockCreateConn + socket.create_connection.side_effect=raise_error + with self.assertRaises(ConnectionError): + KafkaConnection(self.config['host'], self.config['port']) + + def test_send__reconnects_on_dirty_conn(self): + + # Dirty the connection + try: + self.conn._raise_connection_error() + except ConnectionError: + pass + + # Now test that sending attempts to reconnect + self.assertEqual(self.MockCreateConn.call_count, 0) + self.conn.send(self.config['request_id'], self.config['payload']) + self.assertEqual(self.MockCreateConn.call_count, 1) + + def test_send__failure_sets_dirty_connection(self): + + def raise_error(*args): + raise socket.error + + assert isinstance(self.conn._sock, mock.Mock) + self.conn._sock.sendall.side_effect=raise_error + try: + self.conn.send(self.config['request_id'], self.config['payload']) + except ConnectionError: + self.assertIsNone(self.conn._sock) + + def test_recv(self): + + self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload']) + + def test_recv__reconnects_on_dirty_conn(self): + + # Dirty the connection + try: + self.conn._raise_connection_error() + except ConnectionError: + pass + + # Now test that recv'ing attempts to reconnect + self.assertEqual(self.MockCreateConn.call_count, 0) + self.conn.recv(self.config['request_id']) + self.assertEqual(self.MockCreateConn.call_count, 1) + + def test_recv__failure_sets_dirty_connection(self): + + def raise_error(*args): + raise socket.error + + # test that recv'ing attempts to reconnect + assert isinstance(self.conn._sock, mock.Mock) + self.conn._sock.recv.side_effect=raise_error + try: + self.conn.recv(self.config['request_id']) + except ConnectionError: + self.assertIsNone(self.conn._sock) + + def test_recv__doesnt_consume_extra_data_in_stream(self): + + # Here just test that each call to recv will return a single payload + self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload']) + self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload2']) + + def test_get_connected_socket(self): + s = self.conn.get_connected_socket() + + self.assertEqual(s, self.MockCreateConn()) + + def test_get_connected_socket_on_dirty_conn(self): + # Dirty the connection + try: + self.conn._raise_connection_error() + except ConnectionError: + pass + + # Test that get_connected_socket tries to connect + self.assertEqual(self.MockCreateConn.call_count, 0) + self.conn.get_connected_socket() + self.assertEqual(self.MockCreateConn.call_count, 1) + + def test_close__object_is_reusable(self): + + # test that sending to a closed connection + # will re-connect and send data to the socket + self.conn.close() + self.conn.send(self.config['request_id'], self.config['payload']) + self.assertEqual(self.MockCreateConn.call_count, 1) + self.conn._sock.sendall.assert_called_with(self.config['payload']) + + +class TestKafkaConnection(unittest.TestCase): + @mock.patch('socket.create_connection') + def test_copy(self, socket): + """KafkaConnection copies work as expected""" + + conn = KafkaConnection('kafka', 9092) + self.assertEqual(socket.call_count, 1) + + copy = conn.copy() + self.assertEqual(socket.call_count, 1) + self.assertEqual(copy.host, 'kafka') + self.assertEqual(copy.port, 9092) + self.assertEqual(copy._sock, None) + + copy.reinit() + self.assertEqual(socket.call_count, 2) + self.assertNotEqual(copy._sock, None) + + @mock.patch('socket.create_connection') + def test_copy_thread(self, socket): + """KafkaConnection copies work in other threads""" + + err = [] + copy = KafkaConnection('kafka', 9092).copy() + + def thread_func(err, copy): + try: + self.assertEqual(copy.host, 'kafka') + self.assertEqual(copy.port, 9092) + self.assertNotEqual(copy._sock, None) + except Exception as e: + err.append(e) + else: + err.append(None) + thread = Thread(target=thread_func, args=(err, copy)) + thread.start() + thread.join() + + self.assertEqual(err, [None]) + self.assertEqual(socket.call_count, 2) |