예제 #1
0
 def test_ctor(self):
     conn = AIOKafkaConnection('localhost', 1234, loop=self.loop)
     self.assertEqual('localhost', conn.host)
     self.assertEqual(1234, conn.port)
     self.assertTrue('KafkaConnection' in conn.__repr__())
     self.assertIsNone(conn._reader)
     self.assertIsNone(conn._writer)
예제 #2
0
 async def test_ctor(self):
     conn = AIOKafkaConnection('localhost', 1234)
     self.assertEqual('localhost', conn.host)
     self.assertEqual(1234, conn.port)
     self.assertTrue('KafkaConnection' in conn.__repr__())
     self.assertIsNone(conn._reader)
     self.assertIsNone(conn._writer)
예제 #3
0
    async def test__do_sasl_handshake_v1(self):
        host, port = self.kafka_host, self.kafka_port

        # setup connection with mocked send and send_bytes
        conn = AIOKafkaConnection(
            host=host, port=port, loop=self.loop,
            sasl_mechanism="PLAIN",
            sasl_plain_username="******",
            sasl_plain_password="******",
            security_protocol="SASL_PLAINTEXT"
        )
        conn.close = close_mock = mock.MagicMock()

        supported_mechanisms = ["PLAIN"]
        error_class = NoError
        auth_error_class = NoError

        async def mock_send(request, expect_response=True):
            if request.API_KEY == SaslHandShakeRequest[0].API_KEY:
                assert request.API_VERSION == 1
                return SaslHandShakeResponse[1](
                    error_code=error_class.errno,
                    enabled_mechanisms=supported_mechanisms
                )
            else:
                assert request.API_KEY == SaslAuthenticateRequest[0].API_KEY
                return SaslAuthenticateResponse[0](
                    error_code=auth_error_class.errno,
                    error_message="",
                    sasl_auth_bytes=b""
                )

        conn.send = mock.Mock(side_effect=mock_send)
        conn._version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 1]
        })

        await conn._do_sasl_handshake()

        supported_mechanisms = ["GSSAPI"]
        with self.assertRaises(UnsupportedSaslMechanismError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
        supported_mechanisms = ["PLAIN"]

        auth_error_class = IllegalSaslStateError
        close_mock.reset()

        with self.assertRaises(IllegalSaslStateError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
        auth_error_class = NoError

        error_class = UnknownError
        close_mock.reset()

        with self.assertRaises(UnknownError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
예제 #4
0
    def test_send_to_closed(self):
        host, port = self.kafka_host, self.kafka_port
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)
        request = MetadataRequest([])
        with self.assertRaises(ConnectionError):
            yield from conn.send(request)

        conn._writer = mock.MagicMock()
        conn._writer.write.side_effect = OSError('mocked writer is closed')

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
예제 #5
0
    async def test_send_to_closed(self):
        host, port = self.kafka_host, self.kafka_port
        conn = AIOKafkaConnection(host=host, port=port)
        request = MetadataRequest([])
        with self.assertRaises(KafkaConnectionError):
            await conn.send(request)

        conn._writer = mock.MagicMock()
        conn._writer.write.side_effect = OSError('mocked writer is closed')

        with self.assertRaises(KafkaConnectionError):
            await conn.send(request)
예제 #6
0
    def test_osserror_in_reader_task(self):
        host, port = self.kafka_host, self.kafka_port

        @asyncio.coroutine
        def invoke_osserror(*a, **kw):
            yield from asyncio.sleep(0.1, loop=self.loop)
            raise OSError('test oserror')

        request = MetadataRequest([])
        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        reader.readexactly.return_value = invoke_osserror()
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio. async (conn._read(), loop=self.loop)

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
        self.assertEqual(conn.connected(), False)
예제 #7
0
    def test_send_to_closed(self):
        host, port = self.kafka_host, self.kafka_port
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)
        request = MetadataRequest([])
        with self.assertRaises(ConnectionError):
            yield from conn.send(request)

        @asyncio.coroutine
        def invoke_osserror(*a, **kw):
            yield from asyncio.sleep(0.1, loop=self.loop)
            raise OSError('mocked writer is closed')
        conn._writer = mock.MagicMock()
        conn._writer.write.side_effect = OSError('mocked writer is closed')

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
예제 #8
0
    def test_send_to_closed(self):
        host, port = self.kafka_host, self.kafka_port
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)
        request = MetadataRequest([])
        with self.assertRaises(ConnectionError):
            yield from conn.send(request)

        @asyncio.coroutine
        def invoke_osserror(*a, **kw):
            yield from asyncio.sleep(0.1, loop=self.loop)
            raise OSError('mocked writer is closed')

        conn._writer = mock.MagicMock()
        conn._writer.write.side_effect = OSError('mocked writer is closed')

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
예제 #9
0
    def test_osserror_in_reader_task(self):
        host, port = self.kafka_host, self.kafka_port

        @asyncio.coroutine
        def invoke_osserror(*a, **kw):
            yield from asyncio.sleep(0.1, loop=self.loop)
            raise OSError('test oserror')

        request = MetadataRequest([])
        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        reader.readexactly.return_value = invoke_osserror()
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio.async(conn._read(), loop=self.loop)

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
        self.assertEqual(conn.connected(), False)
예제 #10
0
    def test_osserror_in_reader_task(self):
        host, port = self.kafka_host, self.kafka_port

        @asyncio.coroutine
        def invoke_osserror(*a, **kw):
            yield from asyncio.sleep(0.1, loop=self.loop)
            raise OSError

        encoder = KafkaProtocol.encode_metadata_request
        request = encoder(client_id=b"aiokafka-python",
                          correlation_id=1,
                          payloads=())

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        reader.readexactly.return_value = invoke_osserror()
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio. async (conn._read(), loop=self.loop)

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
예제 #11
0
    def test__do_sasl_handshake_v1(self):
        host, port = self.kafka_host, self.kafka_port

        # setup connection with mocked send and send_bytes
        conn = AIOKafkaConnection(
            host=host, port=port, loop=self.loop,
            sasl_mechanism="PLAIN",
            sasl_plain_username="******",
            sasl_plain_password="******",
            security_protocol="SASL_PLAINTEXT"
        )
        conn.close = close_mock = mock.MagicMock()

        supported_mechanisms = ["PLAIN"]
        error_class = NoError
        auth_error_class = NoError

        @asyncio.coroutine
        def mock_send(request, expect_response=True):
            if request.API_KEY == SaslHandShakeRequest[0].API_KEY:
                assert request.API_VERSION == 1
                return SaslHandShakeResponse[1](
                    error_code=error_class.errno,
                    enabled_mechanisms=supported_mechanisms
                )
            else:
                assert request.API_KEY == SaslAuthenticateRequest[0].API_KEY
                return SaslAuthenticateResponse[0](
                    error_code=auth_error_class.errno,
                    error_message="",
                    sasl_auth_bytes=b""
                )

        conn.send = mock.Mock(side_effect=mock_send)
        conn._version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 1]
        })

        yield from conn._do_sasl_handshake()

        supported_mechanisms = ["GSSAPI"]
        with self.assertRaises(UnsupportedSaslMechanismError):
            yield from conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
        supported_mechanisms = ["PLAIN"]

        auth_error_class = IllegalSaslStateError
        close_mock.reset()

        with self.assertRaises(IllegalSaslStateError):
            yield from conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
        auth_error_class = NoError

        error_class = UnknownError
        close_mock.reset()

        with self.assertRaises(UnknownError):
            yield from conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
예제 #12
0
    def test_invalid_correlation_id(self):
        host, port = self.kafka_host, self.kafka_port

        request = MetadataRequest([])

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = MetadataResponse(brokers=[], topics=[]).encode()
        resp = int32.pack(999) + resp  # set invalid correlation id
        reader.readexactly.side_effect = [
            asyncio.coroutine(lambda *a, **kw: int32.pack(len(resp)))(),
            asyncio.coroutine(lambda *a, **kw: resp)()
        ]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio. async (conn._read(), loop=self.loop)

        with self.assertRaises(CorrelationIdError):
            yield from conn.send(request)
예제 #13
0
    def test_correlation_id_on_group_coordinator_req(self):
        host, port = self.kafka_host, self.kafka_port

        request = GroupCoordinatorRequest(consumer_group='test')

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = GroupCoordinatorResponse(error_code=0,
                                        coordinator_id=22,
                                        host='127.0.0.1',
                                        port=3333).encode()
        resp = int32.pack(0) + resp  # set correlation id to 0
        reader.readexactly.side_effect = [
            asyncio.coroutine(lambda *a, **kw: int32.pack(len(resp)))(),
            asyncio.coroutine(lambda *a, **kw: resp)()
        ]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio. async (conn._read(), loop=self.loop)

        response = yield from conn.send(request)
        self.assertIsInstance(response, GroupCoordinatorResponse)
        self.assertEqual(response.error_code, 0)
        self.assertEqual(response.coordinator_id, 22)
        self.assertEqual(response.host, '127.0.0.1')
        self.assertEqual(response.port, 3333)
예제 #14
0
    async def test__do_sasl_handshake_v0(self):
        host, port = self.kafka_host, self.kafka_port

        # setup connection with mocked send and send_bytes
        conn = AIOKafkaConnection(
            host=host, port=port, loop=self.loop,
            sasl_mechanism="PLAIN",
            sasl_plain_username="******",
            sasl_plain_password="******"
        )
        conn.close = close_mock = mock.MagicMock()

        supported_mechanisms = ["PLAIN"]
        error_class = NoError

        async def mock_send(request, expect_response=True):
            return SaslHandShakeResponse[0](
                error_code=error_class.errno,
                enabled_mechanisms=supported_mechanisms
            )

        async def mock_sasl_send(payload, expect_response):
            return b""

        conn.send = mock.Mock(side_effect=mock_send)
        conn._send_sasl_token = mock.Mock(side_effect=mock_sasl_send)
        conn._version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 0]
        })

        await conn._do_sasl_handshake()

        supported_mechanisms = ["GSSAPI"]
        with self.assertRaises(UnsupportedSaslMechanismError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)

        error_class = UnknownError
        close_mock.reset()

        with self.assertRaises(UnknownError):
            await conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
예제 #15
0
    def test__send_sasl_token(self):
        # Before Kafka 1.0.0 SASL was performed on the wire without
        # KAFKA_HEADER in the protocol. So we needed another private
        # function to send `raw` data with only length prefixed

        # setup connection with mocked transport and protocol
        conn = AIOKafkaConnection(
            host="", port=9999, loop=self.loop
        )
        conn.close = mock.MagicMock()
        conn._writer = mock.MagicMock()
        out_buffer = []
        conn._writer.write = mock.Mock(side_effect=out_buffer.append)
        conn._reader = mock.MagicMock()
        self.assertEqual(len(conn._requests), 0)

        # Successful send
        fut = conn._send_sasl_token(b"Super data")
        self.assertEqual(b''.join(out_buffer), b"\x00\x00\x00\nSuper data")
        self.assertEqual(len(conn._requests), 1)
        out_buffer.clear()

        # Resolve the request
        conn._requests[0][2].set_result(None)
        conn._requests.clear()
        yield from fut

        # Broken pipe error
        conn._writer.write.side_effect = OSError
        with self.assertRaises(ConnectionError):
            conn._send_sasl_token(b"Super data")
        self.assertEqual(out_buffer, [])
        self.assertEqual(len(conn._requests), 0)
        self.assertEqual(conn.close.call_count, 1)

        conn._writer = None
        with self.assertRaises(ConnectionError):
            conn._send_sasl_token(b"Super data")
        # We don't need to close 2ce
        self.assertEqual(conn.close.call_count, 1)
예제 #16
0
    def test_correlation_id_on_group_coordinator_req(self):
        host, port = self.kafka_host, self.kafka_port

        request = GroupCoordinatorRequest(consumer_group='test')

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = GroupCoordinatorResponse(
            error_code=0, coordinator_id=22,
            host='127.0.0.1', port=3333).encode()
        resp = int32.pack(0) + resp  # set correlation id to 0
        reader.readexactly.side_effect = [
            asyncio.coroutine(lambda *a, **kw: int32.pack(len(resp)))(),
            asyncio.coroutine(lambda *a, **kw: resp)()]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio.async(conn._read(), loop=self.loop)

        response = yield from conn.send(request)
        self.assertIsInstance(response, GroupCoordinatorResponse)
        self.assertEqual(response.error_code, 0)
        self.assertEqual(response.coordinator_id, 22)
        self.assertEqual(response.host, '127.0.0.1')
        self.assertEqual(response.port, 3333)
예제 #17
0
    def test_osserror_in_reader_task(self):
        host, port = self.server.host, self.server.port

        @asyncio.coroutine
        def invoke_osserror(*a, **kw):
            yield from asyncio.sleep(0.1, loop=self.loop)
            raise OSError

        encoder = KafkaProtocol.encode_metadata_request
        request = encoder(client_id=b"aiokafka-python", correlation_id=1, payloads=())

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        reader.readexactly.return_value = invoke_osserror()
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio.async(conn._read(), loop=self.loop)

        with self.assertRaises(ConnectionError):
            yield from conn.send(request)
예제 #18
0
    async def test_invalid_correlation_id(self):
        host, port = self.kafka_host, self.kafka_port

        request = MetadataRequest([])

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = MetadataResponse(brokers=[], topics=[])
        resp = resp.encode()
        resp = int32.pack(999) + resp  # set invalid correlation id

        async def first_resp(*args: Any, **kw: Any):
            return int32.pack(len(resp))

        async def second_resp(*args: Any, **kw: Any):
            return resp

        reader.readexactly.side_effect = [first_resp(), second_resp()]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = conn._create_reader_task()

        with self.assertRaises(CorrelationIdError):
            await conn.send(request)
예제 #19
0
    def test_invalid_correlation_id(self):
        host, port = self.kafka_host, self.kafka_port

        request = MetadataRequest([])

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port, loop=self.loop)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = MetadataResponse(brokers=[], topics=[]).encode()
        resp = int32.pack(999) + resp  # set invalid correlation id
        reader.readexactly.side_effect = [
            asyncio.coroutine(lambda *a, **kw: int32.pack(len(resp)))(),
            asyncio.coroutine(lambda *a, **kw: resp)()]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = asyncio.async(conn._read(), loop=self.loop)

        with self.assertRaises(CorrelationIdError):
            yield from conn.send(request)
예제 #20
0
    def test__do_sasl_handshake_v0(self):
        host, port = self.kafka_host, self.kafka_port

        # setup connection with mocked send and send_bytes
        conn = AIOKafkaConnection(
            host=host, port=port, loop=self.loop,
            sasl_mechanism="PLAIN",
            sasl_plain_username="******",
            sasl_plain_password="******"
        )
        conn.close = close_mock = mock.MagicMock()

        supported_mechanisms = ["PLAIN"]
        error_class = NoError

        @asyncio.coroutine
        def mock_send(request, expect_response=True):
            return SaslHandShakeResponse[0](
                error_code=error_class.errno,
                enabled_mechanisms=supported_mechanisms
            )

        @asyncio.coroutine
        def mock_sasl_send(payload, expect_response):
            return b""

        conn.send = mock.Mock(side_effect=mock_send)
        conn._send_sasl_token = mock.Mock(side_effect=mock_sasl_send)
        conn._version_info = VersionInfo({
            SaslHandShakeRequest[0].API_KEY: [0, 0]
        })

        yield from conn._do_sasl_handshake()

        supported_mechanisms = ["GSSAPI"]
        with self.assertRaises(UnsupportedSaslMechanismError):
            yield from conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)

        error_class = UnknownError
        close_mock.reset()

        with self.assertRaises(UnknownError):
            yield from conn._do_sasl_handshake()
        self.assertTrue(close_mock.call_count)
예제 #21
0
    async def test_correlation_id_on_group_coordinator_req(self):
        host, port = self.kafka_host, self.kafka_port

        request = GroupCoordinatorRequest(consumer_group='test')

        # setup connection with mocked reader and writer
        conn = AIOKafkaConnection(host=host, port=port)

        # setup reader
        reader = mock.MagicMock()
        int32 = struct.Struct('>i')
        resp = GroupCoordinatorResponse(error_code=0,
                                        coordinator_id=22,
                                        host='127.0.0.1',
                                        port=3333)
        resp = resp.encode()
        resp = int32.pack(0) + resp  # set correlation id to 0

        async def first_resp(*args: Any, **kw: Any):
            return int32.pack(len(resp))

        async def second_resp(*args: Any, **kw: Any):
            return resp

        reader.readexactly.side_effect = [first_resp(), second_resp()]
        writer = mock.MagicMock()

        conn._reader = reader
        conn._writer = writer
        # invoke reader task
        conn._read_task = conn._create_reader_task()

        response = await conn.send(request)
        self.assertIsInstance(response, GroupCoordinatorResponse)
        self.assertEqual(response.error_code, 0)
        self.assertEqual(response.coordinator_id, 22)
        self.assertEqual(response.host, '127.0.0.1')
        self.assertEqual(response.port, 3333)
예제 #22
0
 def setUp(self):
     self.loop = asyncio.new_event_loop()
     asyncio.set_event_loop(None)
     self.conn = AIOKafkaConnection('localhost', 1234, loop=self.loop)
예제 #23
0
    async def test__send_sasl_token(self):
        # Before Kafka 1.0.0 SASL was performed on the wire without
        # KAFKA_HEADER in the protocol. So we needed another private
        # function to send `raw` data with only length prefixed

        # setup connection with mocked transport and protocol
        conn = AIOKafkaConnection(host="", port=9999)
        conn.close = mock.MagicMock()
        conn._writer = mock.MagicMock()
        out_buffer = []
        conn._writer.write = mock.Mock(side_effect=out_buffer.append)
        conn._reader = mock.MagicMock()
        self.assertEqual(len(conn._requests), 0)

        # Successful send
        fut = conn._send_sasl_token(b"Super data")
        self.assertEqual(b''.join(out_buffer), b"\x00\x00\x00\nSuper data")
        self.assertEqual(len(conn._requests), 1)
        out_buffer.clear()

        # Resolve the request
        conn._requests[0][2].set_result(None)
        conn._requests.clear()
        await fut

        # Broken pipe error
        conn._writer.write.side_effect = OSError
        with self.assertRaises(KafkaConnectionError):
            conn._send_sasl_token(b"Super data")
        self.assertEqual(out_buffer, [])
        self.assertEqual(len(conn._requests), 0)
        self.assertEqual(conn.close.call_count, 1)

        conn._writer = None
        with self.assertRaises(KafkaConnectionError):
            conn._send_sasl_token(b"Super data")
        # We don't need to close 2ce
        self.assertEqual(conn.close.call_count, 1)