import asyncio
import unittest
from unittest import mock
from aiokafka.client import AIOKafkaClient
from kafka.common import KafkaUnavailableError


class TestAIOKafkaClinet(unittest.TestCase):

    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(None)

    def tearDown(self):
        self.loop.close()

    def test_init_with_list(self):
        client = AIOKafkaClient(
            ['kafka01:9092', 'kafka02:9092', 'kafka03:9092'], loop=self.loop)

        self.assertEqual(sorted({'kafka01': 9092,
                                 'kafka02': 9092,
                                 'kafka03': 9092}.items()),
                         sorted(client.hosts))

    def test_init_with_csv(self):
        client = AIOKafkaClient('kafka01:9092,kafka02:9092,kafka03:9092',
                                loop=self.loop)

        self.assertEqual(sorted({'kafka01': 9092,
                                 'kafka02': 9092,
                                 'kafka03': 9092}.items()),
                         sorted(client.hosts))

    def test_send_broker_unaware_request_fail(self):
        'Tests that call fails when all hosts are unavailable'

        mocked_conns = {
            ('kafka01', 9092): mock.MagicMock(),
            ('kafka02', 9092): mock.MagicMock()
        }

        # inject KafkaConnection side effects
        mocked_conns[('kafka01', 9092)].send.side_effect = RuntimeError(
            "kafka01 went away (unittest)")
        mocked_conns[('kafka02', 9092)].send.side_effect = RuntimeError(
            "Kafka02 went away (unittest)")

        client = AIOKafkaClient(['kafka01:9092', 'kafka02:9092'],
                                loop=self.loop)
        client._conns = mocked_conns

        @asyncio.coroutine
        def go():
            with self.assertRaises(KafkaUnavailableError):
                yield from client._send_broker_unaware_request(
                    payloads=['fake request'],
                    encoder_fn=mock.MagicMock(
                        return_value=b'fake encoded message'),
                    decoder_fn=lambda x: x)

            for key, conn in mocked_conns.items():
                conn.send.assert_called_with(mock.ANY,
                                             b'fake encoded message')

        self.loop.run_until_complete(go())

@@ -2336,8 +2336,1259 @@
 e(go())%0A
+%0A    def test_send_broker_unaware_request(self):%0A        'Tests that call works when at least one of the host is available'%0A%0A        mocked_conns = %7B%0A            ('kafka01', 9092): mock.MagicMock(),%0A            ('kafka02', 9092): mock.MagicMock(),%0A            ('kafka03', 9092): mock.MagicMock()%0A        %7D%0A        # inject KafkaConnection side effects%0A        mocked_conns%5B('kafka01', 9092)%5D.send.side_effect = RuntimeError(%0A            %22kafka01 went away (unittest)%22)%0A%0A        @asyncio.coroutine%0A        def recv(request_id):%0A            return b'valid response'%0A%0A        mocked_conns%5B('kafka02', 9092)%5D.recv.side_effect = recv%0A        mocked_conns%5B('kafka03', 9092)%5D.send.side_effect = RuntimeError(%0A            %22kafka03 went away (unittest)%22)%0A%0A        client = AIOKafkaClient('kafka01:9092,kafka02:9092', loop=self.loop)%0A        client._conns = mocked_conns%0A%0A        resp = self.loop.run_until_complete(%0A            client._send_broker_unaware_request(payloads=%5Bb'fake request'%5D,%0A                                                encoder_fn=mock.MagicMock(),%0A                                                decoder_fn=lambda x: x))%0A%0A        self.assertEqual(b'valid response', resp)%0A        mocked_conns%5B('kafka02', 9092)%5D.recv.assert_called_with(1)%0A

Add test