Python AIOKafkaClient._get_conn方法代碼示例

# 需要導入模塊: from aiokafka.client import AIOKafkaClient [as 別名]
# 或者: from aiokafka.client.AIOKafkaClient import _get_conn [as 別名]
    def test_client_receive_zero_brokers(self):
        brokers = [
            (0, 'broker_1', 4567),
            (1, 'broker_2', 5678)
        ]
        correct_meta = MetadataResponse(brokers, [])
        bad_response = MetadataResponse([], [])

        @asyncio.coroutine
        def send(*args, **kwargs):
            return bad_response

        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers=['broker_1:4567'],
                                api_version="0.10")
        conn = mock.Mock()
        client._conns = [mock.Mock()]
        client._get_conn = mock.Mock()
        client._get_conn.side_effect = asyncio.coroutine(lambda x: conn)
        conn.send = mock.Mock()
        conn.send.side_effect = send
        client.cluster.update_metadata(correct_meta)
        brokers_before = client.cluster.brokers()

        yield from client._metadata_update(client.cluster, [])

        # There broker list should not be purged
        self.assertNotEqual(client.cluster.brokers(), set([]))
        self.assertEqual(client.cluster.brokers(), brokers_before)

# 需要導入模塊: from aiokafka.client import AIOKafkaClient [as 別名]
# 或者: from aiokafka.client.AIOKafkaClient import _get_conn [as 別名]
    def test_different_connections_in_conn_groups(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers=self.hosts,
            metadata_max_age_ms=10000)
        yield from client.bootstrap()
        self.add_cleanup(client.close)

        node_id = client.get_random_node()
        broker = client.cluster.broker_metadata(node_id)
        client.cluster.add_coordinator(
            node_id, broker.host, broker.port, rack=None,
            purpose=(CoordinationType.GROUP, ""))

        conn1 = yield from client._get_conn(node_id)
        conn2 = yield from client._get_conn(
            node_id, group=ConnectionGroup.COORDINATION)

        self.assertTrue(conn1 is not conn2)
        self.assertEqual((conn1.host, conn1.port), (conn2.host, conn2.port))

# 需要導入模塊: from aiokafka.client import AIOKafkaClient [as 別名]
# 或者: from aiokafka.client.AIOKafkaClient import _get_conn [as 別名]
    def test_send_timeout_deletes_connection(self):
        correct_response = MetadataResponse([], [])

        @asyncio.coroutine
        def send_exception(*args, **kwargs):
            raise asyncio.TimeoutError()

        @asyncio.coroutine
        def send(*args, **kwargs):
            return correct_response

        @asyncio.coroutine
        def get_conn(self, node_id, *, group=0):
            conn_id = (node_id, group)
            if conn_id in self._conns:
                conn = self._conns[conn_id]
                if not conn.connected():
                    del self._conns[conn_id]
                else:
                    return conn

            conn = mock.MagicMock()
            conn.send.side_effect = send
            self._conns[conn_id] = conn
            return conn

        node_id = 0
        conn = mock.MagicMock()
        conn.send.side_effect = send_exception
        conn.connected.return_value = True
        mocked_conns = {(node_id, 0): conn}
        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers=['broker_1:4567'])
        client._conns = mocked_conns
        client._get_conn = types.MethodType(get_conn, client)

        # first send timeouts
        with self.assertRaises(RequestTimedOutError):
            yield from client.send(0, MetadataRequest([]))

        conn.close.assert_called_once_with(
            reason=CloseReason.CONNECTION_TIMEOUT)
        # this happens because conn was closed
        conn.connected.return_value = False

        # second send gets new connection and obtains result
        response = yield from client.send(0, MetadataRequest([]))
        self.assertEqual(response, correct_response)
        self.assertNotEqual(conn, client._conns[(node_id, 0)])

# 需要導入模塊: from aiokafka.client import AIOKafkaClient [as 別名]
# 或者: from aiokafka.client.AIOKafkaClient import _get_conn [as 別名]
    def test_check_version(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        ver = yield from client.check_version()
        self.assertTrue('0.' in ver)
        yield from self.wait_topic(client, 'some_test_topic')
        ver2 = yield from client.check_version()
        self.assertEqual(ver, ver2)
        ver2 = yield from client.check_version(client.get_random_node())
        self.assertEqual(ver, ver2)

        with mock.patch.object(
                AIOKafkaConnection, 'send') as mocked:
            mocked.side_effect = KafkaError('mocked exception')
            with self.assertRaises(UnrecognizedBrokerVersion):
                yield from client.check_version(client.get_random_node())

        client._get_conn = asyncio.coroutine(lambda _: None)
        with self.assertRaises(ConnectionError):
            yield from client.check_version()

# 需要導入模塊: from aiokafka.client import AIOKafkaClient [as 別名]
# 或者: from aiokafka.client.AIOKafkaClient import _get_conn [as 別名]
    def test_check_version(self):
        kafka_version = tuple(int(x) for x in self.kafka_version.split("."))

        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        ver = yield from client.check_version()
        self.assertEqual(kafka_version[:2], ver[:2])
        yield from self.wait_topic(client, 'some_test_topic')
        ver2 = yield from client.check_version()
        self.assertEqual(ver, ver2)
        ver2 = yield from client.check_version(client.get_random_node())
        self.assertEqual(ver, ver2)

        with mock.patch.object(
                AIOKafkaConnection, 'send') as mocked:
            mocked.side_effect = KafkaError('mocked exception')
            with self.assertRaises(UnrecognizedBrokerVersion):
                yield from client.check_version(client.get_random_node())

        client._get_conn = asyncio.coroutine(lambda _, **kw: None)
        with self.assertRaises(ConnectionError):
            yield from client.check_version()
        yield from client.close()

# 需要導入模塊: from aiokafka.client import AIOKafkaClient [as 別名]
# 或者: from aiokafka.client.AIOKafkaClient import _get_conn [as 別名]
    def test_metadata_updated_on_socket_disconnect(self):
        # Related to issue 176. A disconnect means that either we lost
        # connection to the node, or we have a node failure. In both cases
        # there's a high probability that Leader distribution will also change.
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers=self.hosts,
            metadata_max_age_ms=10000)
        yield from client.bootstrap()
        self.add_cleanup(client.close)

        # Init a clonnection
        node_id = client.get_random_node()
        assert node_id is not None
        req = MetadataRequest([])
        yield from client.send(node_id, req)

        # No metadata update pending atm
        self.assertFalse(client._md_update_waiter.done())

        # Connection disconnect should trigger an update
        conn = yield from client._get_conn(node_id)
        conn.close(reason=CloseReason.CONNECTION_BROKEN)
        self.assertTrue(client._md_update_waiter.done())