Python AIOKafkaClient._conns方法代码示例

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    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
        fut1 = asyncio.Future(loop=self.loop)
        fut1.set_exception(RuntimeError("kafka01 went away (unittest)"))
        mocked_conns[("kafka01", 9092)].send.return_value = fut1

        fut2 = asyncio.Future(loop=self.loop)
        fut2.set_exception(RuntimeError("kafka02 went away (unittest)"))
        mocked_conns[("kafka02", 9092)].send.return_value = fut2

        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(b"fake encoded message")

        self.loop.run_until_complete(go())

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_send_broker_unaware_request(self):
        "Tests that call works when at least one of the host is available"

        mocked_conns = {
            ("kafka01", 9092): mock.MagicMock(),
            ("kafka02", 9092): mock.MagicMock(),
            ("kafka03", 9092): mock.MagicMock(),
        }
        # inject KafkaConnection side effects
        fut = asyncio.Future(loop=self.loop)
        fut.set_exception(RuntimeError("kafka01 went away (unittest)"))
        mocked_conns[("kafka01", 9092)].send.return_value = fut

        fut2 = asyncio.Future(loop=self.loop)
        fut2.set_result(b"valid response")
        mocked_conns[("kafka02", 9092)].send.return_value = fut2

        fut3 = asyncio.Future(loop=self.loop)
        fut3.set_exception(RuntimeError("kafka03 went away (unittest)"))
        mocked_conns[("kafka03", 9092)].send.return_value = fut3

        client = AIOKafkaClient("kafka01:9092,kafka02:9092", loop=self.loop)
        client._conns = mocked_conns

        resp = self.loop.run_until_complete(
            client._send_broker_unaware_request(
                payloads=[b"fake request"], encoder_fn=mock.MagicMock(), decoder_fn=lambda x: x
            )
        )

        self.assertEqual(b"valid response", resp)

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_ensure_topic_exists(self, protocol):

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        fut = asyncio.Future(loop=self.loop)
        fut.set_result(b"response")
        mocked_conns[("broker_1", 4567)].send.return_value = fut
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [
            TopicMetadata("topic_still_creating", NO_LEADER, []),
            TopicMetadata("topic_doesnt_exist", UNKNOWN_TOPIC_OR_PARTITION, []),
            TopicMetadata(
                "topic_noleaders",
                NO_ERROR,
                [
                    PartitionMetadata("topic_noleaders", 0, -1, [], [], NO_LEADER),
                    PartitionMetadata("topic_noleaders", 1, -1, [], [], NO_LEADER),
                ],
            ),
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())

        with self.assertRaises(UnknownTopicOrPartitionError):
            self.loop.run_until_complete(client.ensure_topic_exists("topic_doesnt_exist", timeout=1))

        with self.assertRaises(KafkaTimeoutError):
            self.loop.run_until_complete(client.ensure_topic_exists("topic_still_creating", timeout=1))

        # This should not raise
        self.loop.run_until_complete(client.ensure_topic_exists("topic_noleaders", timeout=1))

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_send_produce_request_raises_when_noleader(self, protocol):
        """Send producer request raises LeaderNotAvailableError
           if leader is not available"""

        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [
            TopicMetadata(
                "topic_noleader",
                NO_ERROR,
                [
                    PartitionMetadata("topic_noleader", 0, -1, [], [], NO_LEADER),
                    PartitionMetadata("topic_noleader", 1, -1, [], [], NO_LEADER),
                ],
            )
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())

        requests = [ProduceRequest("topic_noleader", 0, [create_message("a"), create_message("b")])]

        with self.assertRaises(LeaderNotAvailableError):
            self.loop.run_until_complete(client.send_produce_request(requests))

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_get_leader_for_unassigned_partitions(self, protocol):
        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [
            TopicMetadata("topic_no_partitions", NO_LEADER, []),
            TopicMetadata("topic_unknown", UNKNOWN_TOPIC_OR_PARTITION, []),
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())

        self.assertDictEqual({}, client._topics_to_brokers)

        with self.assertRaises(LeaderNotAvailableError):
            self.loop.run_until_complete(client._get_leader_for_partition("topic_no_partitions", 0))

        with self.assertRaises(UnknownTopicOrPartitionError):
            self.loop.run_until_complete(client._get_leader_for_partition("topic_unknown", 0))

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_has_metadata_for_topic(self, protocol):
        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [
            TopicMetadata("topic_still_creating", NO_LEADER, []),
            TopicMetadata("topic_doesnt_exist", UNKNOWN_TOPIC_OR_PARTITION, []),
            TopicMetadata(
                "topic_noleaders",
                NO_ERROR,
                [
                    PartitionMetadata("topic_noleaders", 0, -1, [], [], NO_LEADER),
                    PartitionMetadata("topic_noleaders", 1, -1, [], [], NO_LEADER),
                ],
            ),
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())

        # Topics with no partitions return False
        self.assertFalse(client.has_metadata_for_topic("topic_still_creating"))
        self.assertFalse(client.has_metadata_for_topic("topic_doesnt_exist"))

        # Topic with partition metadata, but no leaders return True
        self.assertTrue(client.has_metadata_for_topic("topic_noleaders"))

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [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 _conns [as 别名]
 def test_close(self):
     client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
     m1 = mock.Mock()
     m2 = mock.Mock()
     client._conns = {("host1", 4567): m1, ("host2", 5678): m2}
     client.close()
     self.assertEqual({}, client._conns)
     m1.close.assert_raises_with()
     m2.close.assert_raises_with()

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_get_leader_exceptions_when_noleader(self, protocol):
        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [
            TopicMetadata(
                "topic_noleader",
                NO_ERROR,
                [
                    PartitionMetadata("topic_noleader", 0, -1, [], [], NO_LEADER),
                    PartitionMetadata("topic_noleader", 1, -1, [], [], NO_LEADER),
                ],
            )
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())
        self.assertDictEqual(
            {TopicAndPartition("topic_noleader", 0): None, TopicAndPartition("topic_noleader", 1): None},
            client._topics_to_brokers,
        )

        # No leader partitions -- raise LeaderNotAvailableError
        with self.assertRaises(LeaderNotAvailableError):
            self.assertIsNone(self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 0)))
        with self.assertRaises(LeaderNotAvailableError):
            self.assertIsNone(self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 1)))

        # Unknown partitions -- raise UnknownTopicOrPartitionError
        with self.assertRaises(UnknownTopicOrPartitionError):
            self.assertIsNone(self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 2)))

        topics = [
            TopicMetadata(
                "topic_noleader",
                NO_ERROR,
                [
                    PartitionMetadata("topic_noleader", 0, 0, [0, 1], [0, 1], NO_ERROR),
                    PartitionMetadata("topic_noleader", 1, 1, [1, 0], [1, 0], NO_ERROR),
                ],
            )
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)
        self.assertEqual(
            brokers[0], self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 0))
        )
        self.assertEqual(
            brokers[1], self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 1))
        )

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [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 _conns [as 别名]
    def test_send_produce_request_raises_when_topic_unknown(self, protocol):
        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [TopicMetadata("topic_doesnt_exist", UNKNOWN_TOPIC_OR_PARTITION, [])]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())

        requests = [ProduceRequest("topic_doesnt_exist", 0, [create_message("a"), create_message("b")])]

        with self.assertRaises(UnknownTopicOrPartitionError):
            self.loop.run_until_complete(client.send_produce_request(requests))

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_get_leader_for_partitions_reloads_metadata(self, protocol):
        "Get leader for partitions reload metadata if it is not available"

        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [TopicMetadata("topic_no_partitions", NO_LEADER, [])]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())

        # topic metadata is loaded but empty
        self.assertDictEqual({}, client._topics_to_brokers)

        topics = [
            TopicMetadata(
                "topic_one_partition",
                NO_ERROR,
                [PartitionMetadata("topic_no_partition", 0, 0, [0, 1], [0, 1], NO_ERROR)],
            )
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        # calling _get_leader_for_partition (from any broker aware request)
        # will try loading metadata again for the same topic
        leader = self.loop.run_until_complete(client._get_leader_for_partition("topic_one_partition", 0))

        self.assertEqual(brokers[0], leader)
        self.assertDictEqual({TopicAndPartition("topic_one_partition", 0): brokers[0]}, client._topics_to_brokers)

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_load_metadata(self):
        brokers = [
            (0, 'broker_1', 4567),
            (1, 'broker_2', 5678)
        ]

        topics = [
            (NO_ERROR, 'topic_1', [
                (NO_ERROR, 0, 1, [1, 2], [1, 2])
            ]),
            (NO_ERROR, 'topic_2', [
                (NO_LEADER, 0, -1, [], []),
                (NO_LEADER, 1, 1, [], []),
            ]),
            (NO_LEADER, 'topic_no_partitions', []),
            (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []),
            (NO_ERROR, 'topic_3', [
                (NO_ERROR, 0, 0, [0, 1], [0, 1]),
                (NO_ERROR, 1, 1, [1, 0], [1, 0]),
                (NO_ERROR, 2, 0, [0, 1], [0, 1])
            ]),
            (NO_ERROR, 'topic_4', [
                (NO_ERROR, 0, 0, [0, 1], [0, 1]),
                (REPLICA_NOT_AVAILABLE, 1, 1, [1, 0], [1, 0]),
            ])
        ]

        @asyncio.coroutine
        def send(request_id):
            return MetadataResponse(brokers, topics)

        mocked_conns = {0: mock.MagicMock()}
        mocked_conns[0].send.side_effect = send
        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers=['broker_1:4567'])
        task = asyncio.async(client._md_synchronizer(), loop=self.loop)
        client._conns = mocked_conns
        client.cluster.update_metadata(MetadataResponse(brokers[:1], []))

        self.loop.run_until_complete(client.force_metadata_update())
        task.cancel()

        md = client.cluster
        c_brokers = md.brokers()
        self.assertEqual(len(c_brokers), 2)
        self.assertEqual(sorted(brokers), sorted(list(c_brokers)))
        c_topics = md.topics()
        self.assertEqual(len(c_topics), 4)
        self.assertEqual(md.partitions_for_topic('topic_1'), set([0]))
        self.assertEqual(md.partitions_for_topic('topic_2'), set([0, 1]))
        self.assertEqual(md.partitions_for_topic('topic_3'), set([0, 1, 2]))
        self.assertEqual(md.partitions_for_topic('topic_4'), set([0, 1]))
        self.assertEqual(
            md.available_partitions_for_topic('topic_2'), set([1]))

        mocked_conns[0].connected.return_value = False
        is_ready = self.loop.run_until_complete(client.ready(0))
        self.assertEqual(is_ready, False)
        is_ready = self.loop.run_until_complete(client.ready(1))
        self.assertEqual(is_ready, False)
        self.assertEqual(mocked_conns, {})

        with self.assertRaises(NodeNotReadyError):
            self.loop.run_until_complete(client.send(0, None))

# 需要导入模块: from aiokafka.client import AIOKafkaClient [as 别名]
# 或者: from aiokafka.client.AIOKafkaClient import _conns [as 别名]
    def test_load_metadata(self, protocol):
        @asyncio.coroutine
        def recv(request_id):
            return b"response"

        mocked_conns = {("broker_1", 4567): mock.MagicMock()}
        mocked_conns[("broker_1", 4567)].recv.side_effect = recv
        client = AIOKafkaClient(["broker_1:4567"], loop=self.loop)
        client._conns = mocked_conns

        brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)]

        topics = [
            TopicMetadata("topic_1", NO_ERROR, [PartitionMetadata("topic_1", 0, 1, [1, 2], [1, 2], NO_ERROR)]),
            TopicMetadata(
                "topic_noleader",
                NO_ERROR,
                [
                    PartitionMetadata("topic_noleader", 0, -1, [], [], NO_LEADER),
                    PartitionMetadata("topic_noleader", 1, -1, [], [], NO_LEADER),
                ],
            ),
            TopicMetadata("topic_no_partitions", NO_LEADER, []),
            TopicMetadata("topic_unknown", UNKNOWN_TOPIC_OR_PARTITION, []),
            TopicMetadata(
                "topic_3",
                NO_ERROR,
                [
                    PartitionMetadata("topic_3", 0, 0, [0, 1], [0, 1], NO_ERROR),
                    PartitionMetadata("topic_3", 1, 1, [1, 0], [1, 0], NO_ERROR),
                    PartitionMetadata("topic_3", 2, 0, [0, 1], [0, 1], NO_ERROR),
                ],
            ),
            TopicMetadata(
                "topic_4",
                NO_ERROR,
                [
                    PartitionMetadata("topic_4", 0, 0, [0, 1], [0, 1], NO_ERROR),
                    PartitionMetadata("topic_4", 1, 1, [1, 0], [1, 0], REPLICA_NOT_AVAILABLE),
                ],
            ),
        ]
        protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics)

        self.loop.run_until_complete(client.load_metadata_for_topics())
        self.assertDictEqual(
            {
                TopicAndPartition("topic_1", 0): brokers[1],
                TopicAndPartition("topic_noleader", 0): None,
                TopicAndPartition("topic_noleader", 1): None,
                TopicAndPartition("topic_3", 0): brokers[0],
                TopicAndPartition("topic_3", 1): brokers[1],
                TopicAndPartition("topic_3", 2): brokers[0],
                TopicAndPartition("topic_4", 0): brokers[0],
                TopicAndPartition("topic_4", 1): brokers[1],
            },
            client._topics_to_brokers,
        )

        # if we ask for metadata explicitly, it should raise errors
        with self.assertRaises(LeaderNotAvailableError):
            self.loop.run_until_complete(client.load_metadata_for_topics("topic_no_partitions"))

        with self.assertRaises(UnknownTopicOrPartitionError):
            self.loop.run_until_complete(client.load_metadata_for_topics("topic_unknown"))

        # This should not raise
        self.loop.run_until_complete(client.load_metadata_for_topics("topic_no_leader"))

        # This should not raise ReplicaNotAvailableError
        self.loop.run_until_complete(client.load_metadata_for_topics("topic_4"))