diff options
Diffstat (limited to 'test/test_client.py')
| -rw-r--r-- | test/test_client.py | 35 |
1 files changed, 19 insertions, 16 deletions
diff --git a/test/test_client.py b/test/test_client.py index 6a80bbb..274655e 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -1,8 +1,9 @@ import socket from time import sleep -import unittest2 from mock import MagicMock, patch +import six +from . import unittest from kafka import KafkaClient from kafka.common import ( @@ -12,34 +13,34 @@ from kafka.common import ( ConnectionError ) from kafka.conn import KafkaConnection -from kafka.protocol import create_message +from kafka.protocol import KafkaProtocol, create_message from test.testutil import Timer -class TestKafkaClient(unittest2.TestCase): +class TestKafkaClient(unittest.TestCase): def test_init_with_list(self): with patch.object(KafkaClient, 'load_metadata_for_topics'): client = KafkaClient(hosts=['kafka01:9092', 'kafka02:9092', 'kafka03:9092']) - self.assertItemsEqual( - [('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)], - client.hosts) + self.assertEqual( + sorted([('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)]), + sorted(client.hosts)) def test_init_with_csv(self): with patch.object(KafkaClient, 'load_metadata_for_topics'): client = KafkaClient(hosts='kafka01:9092,kafka02:9092,kafka03:9092') - self.assertItemsEqual( - [('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)], - client.hosts) + self.assertEqual( + sorted([('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)]), + sorted(client.hosts)) def test_init_with_unicode_csv(self): with patch.object(KafkaClient, 'load_metadata_for_topics'): client = KafkaClient(hosts=u'kafka01:9092,kafka02:9092,kafka03:9092') - self.assertItemsEqual( - [('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)], - client.hosts) + self.assertEqual( + sorted([('kafka01', 9092), ('kafka02', 9092), ('kafka03', 9092)]), + sorted(client.hosts)) def test_send_broker_unaware_request_fail(self): 'Tests that call fails when all hosts are unavailable' @@ -61,11 +62,12 @@ class TestKafkaClient(unittest2.TestCase): with patch.object(KafkaClient, '_get_conn', side_effect=mock_get_conn): client = KafkaClient(hosts=['kafka01:9092', 'kafka02:9092']) + req = KafkaProtocol.encode_metadata_request(b'client', 0) with self.assertRaises(KafkaUnavailableError): - client._send_broker_unaware_request(1, 'fake request') + client._send_broker_unaware_request(1, req) - for key, conn in mocked_conns.iteritems(): - conn.send.assert_called_with(1, 'fake request') + for key, conn in six.iteritems(mocked_conns): + conn.send.assert_called_with(1, req) def test_send_broker_unaware_request(self): 'Tests that call works when at least one of the host is available' @@ -88,7 +90,8 @@ class TestKafkaClient(unittest2.TestCase): with patch.object(KafkaClient, '_get_conn', side_effect=mock_get_conn): client = KafkaClient(hosts='kafka01:9092,kafka02:9092') - resp = client._send_broker_unaware_request(1, 'fake request') + req = KafkaProtocol.encode_metadata_request(b'client', 0) + resp = client._send_broker_unaware_request(1, req) self.assertEqual('valid response', resp) mocked_conns[('kafka02', 9092)].recv.assert_called_with(1) |