class KafkaIntegrationTestCase(unittest.TestCase):

    topic = None

    @contextmanager
    def silence_loop_exception_handler(self):
        if hasattr(self.loop, "get_exception_handler"):
            orig_handler = self.loop.get_exception_handler()
        else:
            orig_handler = None  # Will set default handler
        self.loop.set_exception_handler(lambda loop, ctx: None)
        yield
        self.loop.set_exception_handler(orig_handler)

    def random_topic_name(self):
        return "topic-{}-{}".format(
            self.id()[self.id().rindex(".") + 1:],
            random_string(10).decode('utf-8')
        )

    def setUp(self):
        super().setUp()
        self._messages = {}
        if not self.topic:
            self.topic = self.random_topic_name()
        self._cleanup = []

    def tearDown(self):
        super().tearDown()
        for coro, args, kw in reversed(self._cleanup):
            task = asyncio.wait_for(coro(*args, **kw), 30)
            self.loop.run_until_complete(task)

    def add_cleanup(self, cb_or_coro, *args, **kw):
        self._cleanup.append((cb_or_coro, args, kw))

    async def wait_topic(self, client, topic):
        client.add_topic(topic)
        for i in range(5):
            ok = await client.force_metadata_update()
            if ok:
                ok = topic in client.cluster.topics()
            if not ok:
                await asyncio.sleep(1)
            else:
                return
        raise AssertionError(f'No topic "{topic}" exists')

    async def send_messages(
        self, partition, messages, *, topic=None,
        timestamp_ms=None, return_inst=False, headers=None
    ):
        topic = topic or self.topic
        ret = []
        producer = AIOKafkaProducer(bootstrap_servers=self.hosts)
        await producer.start()
        try:
            await self.wait_topic(producer.client, topic)

            for msg in messages:
                if isinstance(msg, str):
                    msg = msg.encode()
                elif isinstance(msg, int):
                    msg = str(msg).encode()
                future = await producer.send(
                    topic, msg, partition=partition,
                    timestamp_ms=timestamp_ms, headers=headers)
                resp = await future
                self.assertEqual(resp.topic, topic)
                self.assertEqual(resp.partition, partition)
                if return_inst:
                    ret.append(resp)
                else:
                    ret.append(msg)
        finally:
            await producer.stop()
        return ret

    def assert_message_count(self, messages, num_messages):
        # Make sure we got them all
        self.assertEqual(len(messages), num_messages)

        # Make sure there are no duplicates
        self.assertTrue(all(messages.count(m) == 1 for m in messages))

    def create_ssl_context(self):
        context = create_ssl_context(
            cafile=str(self.ssl_folder / "ca-cert"),
            certfile=str(self.ssl_folder / "cl_client.pem"),
            keyfile=str(self.ssl_folder / "cl_client.key"),
            password="abcdefgh")
        context.check_hostname = False
        return context def run_until_complete(fun):
    assert inspect.iscoroutinefunction(fun), (
        "Can not decorate ordinary function, only async ones"
    )

    @wraps(fun)
    def wrapper(test, *args, **kw):
        loop = test.loop
        timeout = getattr(test, "TEST_TIMEOUT", 120)
        ret = loop.run_until_complete(
            asyncio.wait_for(fun(test, *args, **kw), timeout))
        return ret
    return wrapper def kafka_versions(*versions):
    # Took from kafka-python

    def version_str_to_list(s):
        return list(map(int, s.split('.')))  # e.g., [0, 8, 1, 1]

    def construct_lambda(s):
        if s[0].isdigit():
            op_str = '='
            v_str = s
        elif s[1].isdigit():
            op_str = s[0]  # ! < > =
            v_str = s[1:]
        elif s[2].isdigit():
            op_str = s[0:2]  # >= <=
            v_str = s[2:]
        else:
            raise ValueError('Unrecognized kafka version / operator: %s' % s)

        op_map = {
            '=': operator.eq,
            '!': operator.ne,
            '>': operator.gt,
            '<': operator.lt,
            '>=': operator.ge,
            '<=': operator.le
        }
        op = op_map[op_str]
        version = version_str_to_list(v_str)
        return lambda a: op(version_str_to_list(a), version)

    validators = map(construct_lambda, versions)

    def kafka_versions(func):
        @wraps(func)
        def wrapper(self, *args, **kw):
            kafka_version = self.kafka_version

            if not kafka_version:
                self.skipTest(
                    "no kafka version found. Is this an integration test?")

            for f in validators:
                if not f(kafka_version):
                    self.skipTest("unsupported kafka version")

            return func(self, *args, **kw)
        return wrapper
    return kafka_versions class Sender:
    """ Background processing abstraction for Producer. By all means just
    separates batch delivery and transaction management from the main Producer
    code
    """

    def __init__(
            self, client, *, acks, txn_manager, message_accumulator,
            retry_backoff_ms, linger_ms, request_timeout_ms):
        self.client = client
        self._txn_manager = txn_manager
        self._acks = acks

        self._message_accumulator = message_accumulator
        self._sender_task = None
        self._in_flight = set()
        self._muted_partitions = set()
        self._coordinators = {}
        self._retry_backoff = retry_backoff_ms / 1000
        self._request_timeout_ms = request_timeout_ms
        self._linger_time = linger_ms / 1000

    async def start(self):
        # If producer is idempotent we need to assure we have PID found
        await self._maybe_wait_for_pid()
        self._sender_task = create_task(self._sender_routine())
        self._sender_task.add_done_callback(self._fail_all)

    def _fail_all(self, task):
        """ Called when sender fails. Will fail all pending batches, as they
        will never be delivered as well as fail transaction
        """
        if task.cancelled():
            return
        task_exception = task.exception()

        if task_exception is not None:
            self._message_accumulator.fail_all(task_exception)
            if self._txn_manager is not None:
                self._txn_manager.fatal_error(task_exception)

    @property
    def sender_task(self):
        return self._sender_task

    async def close(self):
        if self._sender_task is not None:
            if not self._sender_task.done():
                self._sender_task.cancel()
                await self._sender_task

    async def _sender_routine(self):
        """ Background task, that sends pending batches to leader nodes for
        batch's partition. This incapsulates same logic as Java's `Sender`
        background thread. Because we use asyncio this is more event based
        loop, rather than counting timeout till next possible even like in
        Java.
        """

        tasks = set()
        txn_task = None  # Track a single task for transaction interactions
        try:
            while True:
                # If indempotence or transactions are turned on we need to
                # have a valid PID to send any request below
                await self._maybe_wait_for_pid()

                waiters = set()
                # As transaction coordination is done via a single, separate
                # socket we do not need to pump it to several nodes, as we do
                # with produce requests.
                # We will only have 1 task at a time and will try to spawn
                # another once that is done.
                txn_manager = self._txn_manager
                muted_partitions = self._muted_partitions
                if txn_manager is not None and \
                        txn_manager.transactional_id is not None:
                    if txn_task is None or txn_task.done():
                        txn_task = self._maybe_do_transactional_request()
                        if txn_task is not None:
                            tasks.add(txn_task)
                        else:
                            # Waiters will not be awaited on exit, tasks will
                            waiters.add(txn_manager.make_task_waiter())
                    # We can't have a race condition between
                    # AddPartitionsToTxnRequest and a ProduceRequest, so we
                    # mute the partition until added.
                    muted_partitions = (
                        muted_partitions | txn_manager.partitions_to_add()
                    )
                batches, unknown_leaders_exist = \
                    self._message_accumulator.drain_by_nodes(
                        ignore_nodes=self._in_flight,
                        muted_partitions=muted_partitions)

                # create produce task for every batch
                for node_id, batches in batches.items():
                    task = create_task(
                        self._send_produce_req(node_id, batches))
                    self._in_flight.add(node_id)
                    for tp in batches:
                        self._muted_partitions.add(tp)
                    tasks.add(task)

                if unknown_leaders_exist:
                    # we have at least one unknown partition's leader,
                    # try to update cluster metadata and wait backoff time
                    fut = self.client.force_metadata_update()
                    waiters |= tasks.union([fut])
                else:
                    fut = self._message_accumulator.data_waiter()
                    waiters |= tasks.union([fut])

                # wait when:
                # * At least one of produce task is finished
                # * Data for new partition arrived
                # * Metadata update if partition leader unknown
                done, _ = await asyncio.wait(
                    waiters,
                    return_when=asyncio.FIRST_COMPLETED)

                # done tasks should never produce errors, if they are it's a
                # bug
                for task in done:
                    task.result()

                tasks -= done

        except asyncio.CancelledError:
            # done tasks should never produce errors, if they are it's a bug
            for task in tasks:
                await task
        except (ProducerFenced, OutOfOrderSequenceNumber,
                TransactionalIdAuthorizationFailed):
            raise
        except Exception:  # pragma: no cover
            log.error("Unexpected error in sender routine", exc_info=True)
            raise KafkaError("Unexpected error during batch delivery")

    async def _maybe_wait_for_pid(self):
        if self._txn_manager is None or self._txn_manager.has_pid():
            return

        while True:
            # If transactions are used we can't just send to a random node, but
            # need to find a suitable coordination node
            if self._txn_manager.transactional_id is not None:
                node_id = await self._find_coordinator(
                    CoordinationType.TRANSACTION,
                    self._txn_manager.transactional_id)
            else:
                node_id = self.client.get_random_node()
            success = await self._do_init_pid(node_id)
            if not success:
                await self.client.force_metadata_update()
            else:
                break

    def _coordinator_dead(self, coordinator_type):
        self._coordinators.pop(coordinator_type, None)

    async def _find_coordinator(self, coordinator_type, coordinator_key):
        assert self._txn_manager is not None
        if coordinator_type in self._coordinators:
            return self._coordinators[coordinator_type]
        while True:
            try:
                coordinator_id = await self.client.coordinator_lookup(
                    coordinator_type, coordinator_key)
            except Errors.TransactionalIdAuthorizationFailed:
                err = Errors.TransactionalIdAuthorizationFailed(
                    self._txn_manager.transactional_id)
                raise err
            except Errors.GroupAuthorizationFailedError:
                err = Errors.GroupAuthorizationFailedError(coordinator_key)
                raise err
            except Errors.CoordinatorNotAvailableError:
                await self.client.force_metadata_update()
                await asyncio.sleep(self._retry_backoff)
                continue
            except Errors.KafkaError as err:
                log.error("FindCoordinator Request failed: %s", err)
                raise KafkaError(repr(err))

            # Try to connect to confirm that the connection can be
            # established.
            ready = await self.client.ready(
                coordinator_id, group=ConnectionGroup.COORDINATION)
            if not ready:
                await asyncio.sleep(self._retry_backoff)
                continue

            self._coordinators[coordinator_type] = coordinator_id

            if coordinator_type == CoordinationType.GROUP:
                log.info(
                    "Discovered coordinator %s for group id %s",
                    coordinator_id,
                    coordinator_key
                )
            else:
                log.info(
                    "Discovered coordinator %s for transactional id %s",
                    coordinator_id,
                    coordinator_key
                )
            return coordinator_id

    async def _do_init_pid(self, node_id):
        handler = InitPIDHandler(self)
        return (await handler.do(node_id))

    ###########################################################################
    # Message delivery handler('s')
    ###########################################################################

    async def _send_produce_req(self, node_id, batches):
        """ Create produce request to node
        If producer configured with `retries`>0 and produce response contain
        "failed" partitions produce request for this partition will try
        resend to broker `retries` times with `retry_timeout_ms` timeouts.

        Arguments:
            node_id (int): kafka broker identifier
            batches (dict): dictionary of {TopicPartition: MessageBatch}
        """
        t0 = time.monotonic()

        handler = SendProduceReqHandler(self, batches)
        await handler.do(node_id)

        # if batches for node is processed in less than a linger seconds
        # then waiting for the remaining time
        sleep_time = self._linger_time - (time.monotonic() - t0)
        if sleep_time > 0:
            await asyncio.sleep(sleep_time)

        self._in_flight.remove(node_id)
        for tp in batches:
            self._muted_partitions.remove(tp)

    ###########################################################################
    # Transaction handler('s')
    ###########################################################################

    def _maybe_do_transactional_request(self):
        txn_manager = self._txn_manager

        # If we have any new partitions, still not added to the transaction
        # we need to do that before committing
        tps = txn_manager.partitions_to_add()
        if tps:
            return create_task(
                self._do_add_partitions_to_txn(tps))

        # We need to add group to transaction before we can commit the offset
        group_id = txn_manager.consumer_group_to_add()
        if group_id is not None:
            return create_task(
                self._do_add_offsets_to_txn(group_id))

        # Now commit the added group's offset
        commit_data = txn_manager.offsets_to_commit()
        if commit_data is not None:
            offsets, group_id = commit_data
            return create_task(
                self._do_txn_offset_commit(offsets, group_id))

        commit_result = txn_manager.needs_transaction_commit()
        if commit_result is not None:
            return create_task(
                self._do_txn_commit(commit_result))

    async def _do_add_partitions_to_txn(self, tps):
        # First assert we have a valid coordinator to send the request to
        node_id = await self._find_coordinator(
            CoordinationType.TRANSACTION, self._txn_manager.transactional_id)
        handler = AddPartitionsToTxnHandler(self, tps)
        return (await handler.do(node_id))

    async def _do_add_offsets_to_txn(self, group_id):
        # First assert we have a valid coordinator to send the request to
        node_id = await self._find_coordinator(
            CoordinationType.TRANSACTION, self._txn_manager.transactional_id)
        handler = AddOffsetsToTxnHandler(self, group_id)
        return (await handler.do(node_id))

    async def _do_txn_offset_commit(self, offsets, group_id):
        # Fast return if nothing to commit
        if not offsets:
            return
        # NOTE: We send this one to GROUP coordinator, not TRANSACTION
        try:
            node_id = await self._find_coordinator(
                CoordinationType.GROUP, group_id)
        except GroupAuthorizationFailedError as exc:
            self._txn_manager.error_transaction(exc)
            return
        log.debug(
            "Sending offset-commit request with %s for group %s to %s",
            offsets, group_id, node_id
        )
        handler = TxnOffsetCommitHandler(self, offsets, group_id)
        return (await handler.do(node_id))

    async def _do_txn_commit(self, commit_result):
        """ Committing transaction should be done with care.
            Transactional requests will be blocked by this coroutine, so no new
        offsets or new partitions will be added.
            Produce requests will be stopped, as accumulator will not be
        yielding any new batches.
        """
        # First we need to ensure that all pending messages were flushed
        # before committing. Note, that this will only flush batches available
        # till this point, no new ones.
        await self._message_accumulator.flush_for_commit()

        txn_manager = self._txn_manager

        # If we never sent any data to begin with, no need to commit
        if txn_manager.is_empty_transaction():
            txn_manager.complete_transaction()
            return

        # First assert we have a valid coordinator to send the request to
        node_id = await self._find_coordinator(
            CoordinationType.TRANSACTION, txn_manager.transactional_id)

        handler = EndTxnHandler(self, commit_result)
        return (await handler.do(node_id)) class InitPIDHandler(BaseHandler):

    def create_request(self):
        txn_manager = self._sender._txn_manager
        return InitProducerIdRequest[0](
            transactional_id=txn_manager.transactional_id,
            transaction_timeout_ms=txn_manager.transaction_timeout_ms)

    def handle_response(self, resp):
        txn_manager = self._sender._txn_manager
        error_type = Errors.for_code(resp.error_code)
        if error_type is Errors.NoError:
            log.debug(
                "Successfully found PID=%s EPOCH=%s for Producer %s",
                resp.producer_id, resp.producer_epoch,
                self._sender.client._client_id)
            self._sender._txn_manager.set_pid_and_epoch(
                resp.producer_id, resp.producer_epoch)
            return
        elif (error_type is CoordinatorNotAvailableError or
                error_type is NotCoordinatorError):
            self._sender._coordinator_dead(CoordinationType.TRANSACTION)
        elif (error_type is CoordinatorLoadInProgressError or
                error_type is ConcurrentTransactions):
            pass
        elif error_type is TransactionalIdAuthorizationFailed:
            raise error_type(txn_manager.transactional_id)
        else:
            log.error(
                "Unexpected error during InitProducerIdRequest: %s",
                error_type)
            raise error_type()

        return self._default_backoff class AddPartitionsToTxnHandler(BaseHandler):
    group = ConnectionGroup.COORDINATION

    def __init__(self, sender, topic_partitions):
        super().__init__(sender)
        self._tps = topic_partitions

    def create_request(self):
        txn_manager = self._sender._txn_manager

        partition_data = collections.defaultdict(list)
        for tp in self._tps:
            partition_data[tp.topic].append(tp.partition)

        req = AddPartitionsToTxnRequest[0](
            transactional_id=txn_manager.transactional_id,
            producer_id=txn_manager.producer_id,
            producer_epoch=txn_manager.producer_epoch,
            topics=list(partition_data.items()))
        return req

    def handle_response(self, resp):
        txn_manager = self._sender._txn_manager

        unauthorized_topics = set()
        for topic, partitions in resp.errors:
            for partition, error_code in partitions:
                tp = TopicPartition(topic, partition)
                error_type = Errors.for_code(error_code)

                if error_type is Errors.NoError:
                    log.debug("Added partition %s to transaction", tp)
                    txn_manager.partition_added(tp)
                elif (error_type is CoordinatorNotAvailableError or
                        error_type is NotCoordinatorError):
                    self._sender._coordinator_dead(
                        CoordinationType.TRANSACTION)
                    return self._default_backoff
                elif error_type is ConcurrentTransactions:
                    # See KAFKA-5477: There is some time between commit and
                    # actual transaction marker write, that will produce this
                    # ConcurrentTransactions. We don't want the 100ms latency
                    # in that case.
                    if not txn_manager.txn_partitions:
                        return BACKOFF_OVERRIDE
                    else:
                        return self._default_backoff
                elif (error_type is CoordinatorLoadInProgressError or
                        error_type is UnknownTopicOrPartitionError):
                    return self._default_backoff
                elif error_type is InvalidProducerEpoch:
                    raise ProducerFenced()
                elif (error_type is InvalidProducerIdMapping or
                        error_type is InvalidTxnState):
                    raise error_type()
                elif error_type is TopicAuthorizationFailedError:
                    unauthorized_topics.add(topic)
                elif error_type is OperationNotAttempted:
                    pass
                elif error_type is TransactionalIdAuthorizationFailed:
                    raise error_type(txn_manager.transactional_id)
                else:
                    log.error(
                        "Could not add partition %s due to unexpected error:"
                        " %s", partition, error_type)
                    raise error_type()
        if unauthorized_topics:
            txn_manager.error_transaction(
                TopicAuthorizationFailedError(unauthorized_topics))
        return class AddOffsetsToTxnHandler(BaseHandler):
    group = ConnectionGroup.COORDINATION

    def __init__(self, sender, group_id):
        super().__init__(sender)
        self._group_id = group_id

    def create_request(self):
        txn_manager = self._sender._txn_manager

        req = AddOffsetsToTxnRequest[0](
            transactional_id=txn_manager.transactional_id,
            producer_id=txn_manager.producer_id,
            producer_epoch=txn_manager.producer_epoch,
            group_id=self._group_id
        )
        return req

    def handle_response(self, resp):
        txn_manager = self._sender._txn_manager
        group_id = self._group_id

        error_type = Errors.for_code(resp.error_code)
        if error_type is Errors.NoError:
            log.debug(
                "Successfully added consumer group %s to transaction", group_id
            )
            txn_manager.consumer_group_added(group_id)
            return
        elif (error_type is CoordinatorNotAvailableError or
                error_type is NotCoordinatorError):
            self._sender._coordinator_dead(CoordinationType.TRANSACTION)
        elif (error_type is CoordinatorLoadInProgressError or
                error_type is ConcurrentTransactions):
            # We will just retry after backoff
            pass
        elif error_type is InvalidProducerEpoch:
            raise ProducerFenced()
        elif error_type is InvalidTxnState:
            raise error_type()
        elif error_type is TransactionalIdAuthorizationFailed:
            raise error_type(txn_manager.transactional_id)
        elif error_type is GroupAuthorizationFailedError:
            txn_manager.error_transaction(error_type(self._group_id))
            return
        else:
            log.error(
                "Could not add consumer group due to unexpected error: %s",
                error_type)
            raise error_type()

        return self._default_backoff class TxnOffsetCommitHandler(BaseHandler):
    group = ConnectionGroup.COORDINATION

    def __init__(self, sender, offsets, group_id):
        super().__init__(sender)
        self._offsets = offsets
        self._group_id = group_id

    def create_request(self):
        txn_manager = self._sender._txn_manager
        # create the offset commit request structure
        offset_data = collections.defaultdict(list)
        for tp, offset in sorted(self._offsets.items()):
            offset_data[tp.topic].append(
                (tp.partition,
                 offset.offset,
                 offset.metadata))

        req = TxnOffsetCommitRequest[0](
            transactional_id=txn_manager.transactional_id,
            group_id=self._group_id,
            producer_id=txn_manager.producer_id,
            producer_epoch=txn_manager.producer_epoch,
            topics=list(offset_data.items())
        )
        return req

    def handle_response(self, resp):
        txn_manager = self._sender._txn_manager
        group_id = self._group_id

        for topic, partitions in resp.errors:
            for partition, error_code in partitions:
                tp = TopicPartition(topic, partition)
                error_type = Errors.for_code(error_code)

                if error_type is Errors.NoError:
                    offset = self._offsets[tp].offset
                    log.debug(
                        "Offset %s for partition %s committed to group %s",
                        offset, tp, group_id)
                    txn_manager.offset_committed(tp, offset, group_id)
                elif (error_type is CoordinatorNotAvailableError or
                        error_type is NotCoordinatorError or
                        # Copied from Java. Not sure why it's only in this case
                        error_type is RequestTimedOutError):
                    self._sender._coordinator_dead(CoordinationType.GROUP)
                    return self._default_backoff
                elif (error_type is CoordinatorLoadInProgressError or
                        error_type is UnknownTopicOrPartitionError):
                    # We will just retry after backoff
                    return self._default_backoff
                elif error_type is InvalidProducerEpoch:
                    raise ProducerFenced()
                elif error_type is TransactionalIdAuthorizationFailed:
                    raise error_type(txn_manager.transactional_id)
                elif error_type is GroupAuthorizationFailedError:
                    exc = error_type(self._group_id)
                    txn_manager.error_transaction(exc)
                    return
                else:
                    log.error(
                        "Could not commit offset for partition %s due to "
                        "unexpected error: %s", partition, error_type)
                    raise error_type() class EndTxnHandler(BaseHandler):
    group = ConnectionGroup.COORDINATION

    def __init__(self, sender, commit_result):
        super().__init__(sender)
        self._commit_result = commit_result

    def create_request(self):
        txn_manager = self._sender._txn_manager
        req = EndTxnRequest[0](
            transactional_id=txn_manager.transactional_id,
            producer_id=txn_manager.producer_id,
            producer_epoch=txn_manager.producer_epoch,
            transaction_result=self._commit_result)
        return req

    def handle_response(self, resp):
        txn_manager = self._sender._txn_manager
        error_type = Errors.for_code(resp.error_code)

        if error_type is Errors.NoError:
            txn_manager.complete_transaction()
            return
        elif (error_type is CoordinatorNotAvailableError or
                error_type is NotCoordinatorError):
            self._sender._coordinator_dead(CoordinationType.TRANSACTION)
        elif (error_type is CoordinatorLoadInProgressError or
                error_type is ConcurrentTransactions):
            # We will just retry after backoff
            pass
        elif error_type is InvalidProducerEpoch:
            raise ProducerFenced()
        elif error_type is InvalidTxnState:
            raise error_type()
        else:
            log.error(
                "Could not end transaction due to unexpected error: %s",
                error_type)
            raise error_type()

        return self._default_backoff class BaseHandler:
    group = ConnectionGroup.DEFAULT

    def __init__(self, sender):
        self._sender = sender
        self._default_backoff = sender._retry_backoff

    async def do(self, node_id):
        req = self.create_request()
        try:
            resp = await self._sender.client.send(
                node_id, req, group=self.group)
        except KafkaError as err:
            log.warning("Could not send %r: %r", req.__class__, err)
            await asyncio.sleep(self._default_backoff)
            return False

        retry_backoff = self.handle_response(resp)
        if retry_backoff is not None:
            await asyncio.sleep(retry_backoff)
            return False  # Failure
        else:
            return True  # Success

    def create_request(self):
        raise NotImplementedError  # pragma: no cover

    def handle_response(self, response):
        raise NotImplementedError  # pragma: no cover class SendProduceReqHandler(BaseHandler):

    def __init__(self, sender, batches):
        super().__init__(sender)
        self._batches = batches
        self._client = sender.client
        self._to_reenqueue = []

    def create_request(self):
        topics = collections.defaultdict(list)
        for tp, batch in self._batches.items():
            topics[tp.topic].append(
                (tp.partition, batch.get_data_buffer())
            )

        if self._client.api_version >= (2, 1):
            version = 7
        elif self._client.api_version >= (2, 0):
            version = 6
        elif self._client.api_version >= (1, 1):
            version = 5
        elif self._client.api_version >= (1, 0):
            version = 4
        elif self._client.api_version >= (0, 11):
            version = 3
        elif self._client.api_version >= (0, 10):
            version = 2
        elif self._client.api_version == (0, 9):
            version = 1
        else:
            version = 0

        kwargs = {}
        if version >= 3:
            if self._sender._txn_manager is not None:
                kwargs['transactional_id'] = \
                    self._sender._txn_manager.transactional_id
            else:
                kwargs['transactional_id'] = None

        request = ProduceRequest[version](
            required_acks=self._sender._acks,
            timeout=self._sender._request_timeout_ms,
            topics=list(topics.items()),
            **kwargs)
        return request

    async def do(self, node_id):
        request = self.create_request()
        try:
            response = await self._client.send(node_id, request)
        except KafkaError as err:
            log.warning(
                "Got error produce response: %s", err)
            if getattr(err, "invalid_metadata", False):
                self._client.force_metadata_update()

            for batch in self._batches.values():
                if not self._can_retry(err, batch):
                    batch.failure(exception=err)
                else:
                    self._to_reenqueue.append(batch)
        else:
            # noacks, just mark batches as "done"
            if request.required_acks == 0:
                for batch in self._batches.values():
                    batch.done_noack()
            else:
                self.handle_response(response)

        if self._to_reenqueue:
            # Wait backoff before reequeue
            await asyncio.sleep(self._default_backoff)

            for batch in self._to_reenqueue:
                self._sender._message_accumulator.reenqueue(batch)
            # If some error started metadata refresh we have to wait before
            # trying again
            await self._client._maybe_wait_metadata()

    def handle_response(self, response):
        for topic, partitions in response.topics:
            for partition_info in partitions:
                global_error = None
                log_start_offset = None
                if response.API_VERSION < 2:
                    partition, error_code, offset = partition_info
                    # Mimic CREATE_TIME to take user provided timestamp
                    timestamp = -1
                elif 2 <= response.API_VERSION <= 4:
                    partition, error_code, offset, timestamp = partition_info
                elif 5 <= response.API_VERSION <= 7:
                    (
                        partition, error_code, offset, timestamp,
                        log_start_offset
                    ) = partition_info
                else:
                    # the ignored parameter is record_error of type
                    # list[(batch_index: int, error_message: str)]
                    (
                        partition, error_code, offset, timestamp,
                        log_start_offset, _, global_error
                    ) = partition_info
                tp = TopicPartition(topic, partition)
                error = Errors.for_code(error_code)
                batch = self._batches.get(tp)
                if batch is None:
                    continue

                if error is Errors.NoError:
                    batch.done(offset, timestamp, log_start_offset)
                elif error is DuplicateSequenceNumber:
                    # If we have received a duplicate sequence error,
                    # it means that the sequence number has advanced
                    # beyond the sequence of the current batch, and we
                    # haven't retained batch metadata on the broker to
                    # return the correct offset and timestamp.
                    #
                    # The only thing we can do is to return success to
                    # the user and not return a valid offset and
                    # timestamp.
                    batch.done(offset, timestamp, log_start_offset)
                elif not self._can_retry(error(), batch):
                    if error is InvalidProducerEpoch:
                        exc = ProducerFenced()
                    elif error is TopicAuthorizationFailedError:
                        exc = error(topic)
                    else:
                        exc = error()
                    batch.failure(exception=exc)
                else:
                    log.warning(
                        "Got error produce response on topic-partition"
                        " %s, retrying. Error: %s", tp, global_error or error)
                    # Ok, we can retry this batch
                    if getattr(error, "invalid_metadata", False):
                        self._client.force_metadata_update()
                    self._to_reenqueue.append(batch)

    def _can_retry(self, error, batch):
        # If indempotence is enabled we never expire batches, but retry until
        # we succeed. We can be sure, that no duplicates will be introduced
        # as long as we set proper sequence, pid and epoch.
        if self._sender._txn_manager is None and batch.expired():
            return False
        # XXX: remove unknown topic check as we fix
        #      https://github.com/dpkp/kafka-python/issues/1155
        if error.retriable or isinstance(error, UnknownTopicOrPartitionError)\
                or error is UnknownTopicOrPartitionError:
            return True
        return False class TransactionManager:

    def __init__(self, transactional_id, transaction_timeout_ms):
        self.transactional_id = transactional_id
        self.transaction_timeout_ms = transaction_timeout_ms
        self.state = TransactionState.UNINITIALIZED

        self._pid_and_epoch = PidAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH)
        self._pid_waiter = create_future()
        self._sequence_numbers = defaultdict(lambda: 0)
        self._transaction_waiter = None
        self._task_waiter = None

        self._txn_partitions = set()
        self._pending_txn_partitions = set()
        self._txn_consumer_group = None
        self._pending_txn_offsets = deque()

    # INDEMPOTANCE PART

    def set_pid_and_epoch(self, pid: int, epoch: int):
        self._pid_and_epoch = PidAndEpoch(pid, epoch)
        self._pid_waiter.set_result(None)
        if self.transactional_id:
            self._transition_to(TransactionState.READY)

    def has_pid(self):
        return self._pid_and_epoch.pid != NO_PRODUCER_ID

    async def wait_for_pid(self):
        if self.has_pid():
            return
        else:
            await self._pid_waiter

    def sequence_number(self, tp: TopicPartition):
        return self._sequence_numbers[tp]

    def increment_sequence_number(self, tp: TopicPartition, increment: int):
        # Java will wrap those automatically, but in Python we will break
        # on `struct.pack` if ints are too big, so we do it here
        seq = self._sequence_numbers[tp] + increment
        if seq > 2 ** 31 - 1:
            seq -= 2 ** 32
        self._sequence_numbers[tp] = seq

    @property
    def producer_id(self):
        return self._pid_and_epoch.pid

    @property
    def producer_epoch(self):
        return self._pid_and_epoch.epoch

    # TRANSACTION PART

    def _transition_to(self, target):
        assert TransactionState.is_transition_valid(self.state, target), \
            f"Invalid state transition {self.state} -> {target}"
        self.state = target

    def begin_transaction(self):
        self._transition_to(TransactionState.IN_TRANSACTION)
        self._transaction_waiter = create_future()

    def committing_transaction(self):
        if self.state == TransactionState.ABORTABLE_ERROR:
            # Raise error to user, we can only abort at this point
            self._transaction_waiter.result()

        self._transition_to(TransactionState.COMMITTING_TRANSACTION)
        self.notify_task_waiter()

    def aborting_transaction(self):
        self._transition_to(TransactionState.ABORTING_TRANSACTION)

        # If we had an abortable error we need to create a new waiter
        if self._transaction_waiter.done():
            self._transaction_waiter = create_future()
        self.notify_task_waiter()

    def complete_transaction(self):
        assert not self._pending_txn_partitions
        assert not self._pending_txn_offsets
        self._transition_to(TransactionState.READY)
        self._txn_partitions.clear()
        self._txn_consumer_group = None
        if not self._transaction_waiter.done():
            self._transaction_waiter.set_result(None)

    def error_transaction(self, exc):
        self._transition_to(TransactionState.ABORTABLE_ERROR)
        self._txn_partitions.clear()
        self._txn_consumer_group = None
        self._pending_txn_partitions.clear()
        for _, _, fut in self._pending_txn_offsets:
            fut.set_exception(exc)
        self._pending_txn_offsets.clear()
        self._transaction_waiter.set_exception(exc)

    def fatal_error(self, exc):
        self._transition_to(TransactionState.FATAL_ERROR)
        self._txn_partitions.clear()
        self._txn_consumer_group = None
        self._pending_txn_partitions.clear()
        for _, _, fut in self._pending_txn_offsets:
            fut.set_exception(exc)
        self._pending_txn_offsets.clear()
        # There may be an abortable error. We just override it
        if self._transaction_waiter.done():
            self._transaction_waiter = create_future()
        self._transaction_waiter.set_exception(exc)

    def maybe_add_partition_to_txn(self, tp: TopicPartition):
        if self.transactional_id is None:
            return
        assert self.is_in_transaction()
        if tp not in self._txn_partitions:
            self._pending_txn_partitions.add(tp)
            self.notify_task_waiter()

    def add_offsets_to_txn(self, offsets, group_id):
        assert self.is_in_transaction()
        assert self.transactional_id
        fut = create_future()
        self._pending_txn_offsets.append(
            (group_id, offsets, fut)
        )
        self.notify_task_waiter()
        return fut

    def is_in_transaction(self):
        return self.state == TransactionState.IN_TRANSACTION

    def partitions_to_add(self):
        return self._pending_txn_partitions

    def consumer_group_to_add(self):
        if self._txn_consumer_group is not None:
            return
        for group_id, _, _ in self._pending_txn_offsets:
            return group_id

    def offsets_to_commit(self):
        if self._txn_consumer_group is None:
            return
        for group_id, offsets, _ in self._pending_txn_offsets:
            return offsets, group_id

    def partition_added(self, tp: TopicPartition):
        self._pending_txn_partitions.remove(tp)
        self._txn_partitions.add(tp)

    def consumer_group_added(self, group_id):
        self._txn_consumer_group = group_id

    def offset_committed(self, tp, offset, group_id):
        pending_group_id, pending_offsets, fut = self._pending_txn_offsets[0]
        assert pending_group_id == group_id
        assert tp in pending_offsets and pending_offsets[tp].offset == offset
        del pending_offsets[tp]

        if not pending_offsets:
            fut.set_result(None)
            self._pending_txn_offsets.popleft()

    @property
    def txn_partitions(self):
        return self._txn_partitions

    def needs_transaction_commit(self):
        if self.state == TransactionState.COMMITTING_TRANSACTION:
            return TransactionResult.COMMIT
        elif self.state == TransactionState.ABORTING_TRANSACTION:
            return TransactionResult.ABORT
        else:
            return

    def is_empty_transaction(self):
        # whether we sent either data to a partition or committed offset
        return (
            len(self.txn_partitions) == 0 and
            self._txn_consumer_group is None
        )

    def is_fatal_error(self):
        return self.state == TransactionState.FATAL_ERROR

    def wait_for_transaction_end(self):
        return self._transaction_waiter

    def notify_task_waiter(self):
        if self._task_waiter is not None and not self._task_waiter.done():
            self._task_waiter.set_result(None)

    def make_task_waiter(self):
        self._task_waiter = create_future()
        return self._task_waiter class TransactionState(Enum):

    UNINITIALIZED = 1
    READY = 2
    IN_TRANSACTION = 3
    COMMITTING_TRANSACTION = 4
    ABORTING_TRANSACTION = 5
    ABORTABLE_ERROR = 6
    FATAL_ERROR = 7

    @classmethod
    def is_transition_valid(cls, source, target):
        if target == cls.READY:
            return source == cls.UNINITIALIZED or \
                source == cls.COMMITTING_TRANSACTION or \
                source == cls.ABORTING_TRANSACTION
        elif target == cls.IN_TRANSACTION:
            return source == cls.READY
        elif target == cls.COMMITTING_TRANSACTION:
            return source == cls.IN_TRANSACTION
        elif target == cls.ABORTING_TRANSACTION:
            return source == cls.IN_TRANSACTION or \
                source == cls.ABORTABLE_ERROR
        elif target == cls.ABORTABLE_ERROR or target == cls.FATAL_ERROR:
            return True class InitProducerIdResponse_v0(Response):
class InitProducerIdRequest_v0(Request):
class AddPartitionsToTxnResponse_v0(Response):
class AddPartitionsToTxnRequest_v0(Request):
class AddOffsetsToTxnResponse_v0(Response):
class AddOffsetsToTxnRequest_v0(Request):
class EndTxnResponse_v0(Response):
class EndTxnRequest_v0(Request):
class TxnOffsetCommitResponse_v0(Response):
class TxnOffsetCommitRequest_v0(Request):
    API_KEY = 22
    API_VERSION = 0
    SCHEMA = Schema(
        ('throttle_time_ms', Int32),
        ('error_code', Int16),
        ('producer_id', Int64),
        ('producer_epoch', Int16),
    )
    API_KEY = 22
    API_VERSION = 0
    RESPONSE_TYPE = InitProducerIdResponse_v0
    SCHEMA = Schema(
        ('transactional_id', String('utf-8')),
        ('transaction_timeout_ms', Int32)
    )
    API_KEY = 24
    API_VERSION = 0
    SCHEMA = Schema(
        ('throttle_time_ms', Int32),
        ('errors', Array(
            ('topic', String('utf-8')),
            ('partition_errors', Array(
                ('partition', Int32),
                ('error_code', Int16)))))
    )
    API_KEY = 24
    API_VERSION = 0
    RESPONSE_TYPE = AddPartitionsToTxnResponse_v0
    SCHEMA = Schema(
        ('transactional_id', String('utf-8')),
        ('producer_id', Int64),
        ('producer_epoch', Int16),
        ('topics', Array(
            ('topic', String('utf-8')),
            ('partitions', Array(Int32))))
    )
    API_KEY = 25
    API_VERSION = 0
    SCHEMA = Schema(
        ('throttle_time_ms', Int32),
        ('error_code', Int16)
    )
    API_KEY = 25
    API_VERSION = 0
    RESPONSE_TYPE = AddOffsetsToTxnResponse_v0
    SCHEMA = Schema(
        ('transactional_id', String('utf-8')),
        ('producer_id', Int64),
        ('producer_epoch', Int16),
        ('group_id', String('utf-8'))
    )
    API_KEY = 26
    API_VERSION = 0
    SCHEMA = Schema(
        ('throttle_time_ms', Int32),
        ('error_code', Int16)
    )
    API_KEY = 26
    API_VERSION = 0
    RESPONSE_TYPE = EndTxnResponse_v0
    SCHEMA = Schema(
        ('transactional_id', String('utf-8')),
        ('producer_id', Int64),
        ('producer_epoch', Int16),
        ('transaction_result', Boolean)
    )
    API_KEY = 28
    API_VERSION = 0
    SCHEMA = Schema(
        ('throttle_time_ms', Int32),
        ('errors', Array(
            ('topic', String('utf-8')),
            ('partition_errors', Array(
                ('partition', Int32),
                ('error_code', Int16)))))
    )
    API_KEY = 28
    API_VERSION = 0
    RESPONSE_TYPE = TxnOffsetCommitResponse_v0
    SCHEMA = Schema(
        ('transactional_id', String('utf-8')),
        ('group_id', String('utf-8')),
        ('producer_id', Int64),
        ('producer_epoch', Int16),
        ('topics', Array(
            ('topic', String('utf-8')),
            ('partitions', Array(
                ('partition', Int32),
                ('offset', Int64),
                ('metadata', String('utf-8'))))))
    ) class MessageAccumulator:
    """Accumulator of messages batched by topic-partition

    Producer adds messages to this accumulator and a background send task
    gets batches per nodes to process it.
    """
    def __init__(
            self, cluster, batch_size, compression_type, batch_ttl, *,
            txn_manager=None, loop=None):
        if loop is None:
            loop = get_running_loop()
        self._loop = loop
        self._batches = collections.defaultdict(collections.deque)
        self._pending_batches = set()
        self._cluster = cluster
        self._batch_size = batch_size
        self._compression_type = compression_type
        self._batch_ttl = batch_ttl
        self._wait_data_future = loop.create_future()
        self._closed = False
        self._api_version = (0, 9)
        self._txn_manager = txn_manager

        self._exception = None  # Critical exception

    def set_api_version(self, api_version):
        self._api_version = api_version

    async def flush(self):
        waiters = []
        for batches in self._batches.values():
            for batch in list(batches):
                waiters.append(batch.future)
        for batch in list(self._pending_batches):
            waiters.append(batch.future)
        if waiters:
            await asyncio.wait(waiters)

    async def flush_for_commit(self):
        waiters = []
        for batches in self._batches.values():
            for batch in batches:
                # We force all buffers to close to finalyze the transaction
                # scope. We should not add anything to this transaction.
                batch._builder.close()
                waiters.append(batch.future)
        for batch in self._pending_batches:
            waiters.append(batch.future)
        # Wait for all waiters to finish. We only wait for the scope we defined
        # above, other batches should not be delivered as part of this
        # transaction
        if waiters:
            await asyncio.wait(waiters)

    def fail_all(self, exception):
        # Close all batches with this exception
        for batches in self._batches.values():
            for batch in batches:
                batch.failure(exception)
        for batch in self._pending_batches:
            batch.failure(exception)
        self._exception = exception

    async def close(self):
        self._closed = True
        await self.flush()

    async def add_message(
        self, tp, key, value, timeout, timestamp_ms=None,
        headers=[]
    ):
        """ Add message to batch by topic-partition
        If batch is already full this method waits (`timeout` seconds maximum)
        until batch is drained by send task
        """
        while True:
            if self._closed:
                # this can happen when producer is closing but try to send some
                # messages in async task
                raise ProducerClosed()
            if self._exception is not None:
                raise copy.copy(self._exception)

            pending_batches = self._batches.get(tp)
            if not pending_batches:
                builder = self.create_builder()
                batch = self._append_batch(builder, tp)
            else:
                batch = pending_batches[-1]

            future = batch.append(key, value, timestamp_ms, headers=headers)
            if future is not None:
                return future
            # Batch is full, can't append data atm,
            # waiting until batch per topic-partition is drained
            start = time.monotonic()
            await batch.wait_drain(timeout)
            timeout -= time.monotonic() - start
            if timeout <= 0:
                raise KafkaTimeoutError()

    def data_waiter(self):
        """ Return waiter future that will be resolved when accumulator contain
        some data for drain
        """
        return self._wait_data_future

    def _pop_batch(self, tp):
        batch = self._batches[tp].popleft()
        not_retry = batch.retry_count == 0
        if self._txn_manager is not None and not_retry:
            assert self._txn_manager.has_pid(), \
                "We should have waited for it in sender routine"
            seq = self._txn_manager.sequence_number(batch.tp)
            self._txn_manager.increment_sequence_number(
                batch.tp, batch.record_count)
            batch.set_producer_state(
                producer_id=self._txn_manager.producer_id,
                producer_epoch=self._txn_manager.producer_epoch,
                base_sequence=seq)
        batch.drain_ready()
        if len(self._batches[tp]) == 0:
            del self._batches[tp]
        self._pending_batches.add(batch)

        if not_retry:
            def cb(fut, batch=batch, self=self):
                self._pending_batches.remove(batch)
            batch.future.add_done_callback(cb)
        return batch

    def reenqueue(self, batch):
        tp = batch.tp
        self._batches[tp].appendleft(batch)
        self._pending_batches.remove(batch)
        batch.reset_drain()

    def drain_by_nodes(self, ignore_nodes, muted_partitions=set()):
        """ Group batches by leader to partition nodes. """
        nodes = collections.defaultdict(dict)
        unknown_leaders_exist = False
        for tp in list(self._batches.keys()):
            # Just ignoring by node is not enough, as leader can change during
            # the cycle
            if tp in muted_partitions:
                continue
            leader = self._cluster.leader_for_partition(tp)
            if leader is None or leader == -1:
                if self._batches[tp][0].expired():
                    # batch is for partition is expired and still no leader,
                    # so set exception for batch and pop it
                    batch = self._pop_batch(tp)
                    if leader is None:
                        err = NotLeaderForPartitionError()
                    else:
                        err = LeaderNotAvailableError()
                    batch.failure(exception=err)
                unknown_leaders_exist = True
                continue
            elif ignore_nodes and leader in ignore_nodes:
                continue

            batch = self._pop_batch(tp)
            # We can get an empty batch here if all `append()` calls failed
            # with validation...
            if not batch.is_empty():
                nodes[leader][tp] = batch
            else:
                # XXX: use something more graceful. We just want to trigger
                # delivery future here, no message futures.
                batch.done_noack()

        # all batches are drained from accumulator
        # so create "wait data" future again for waiting new data in send
        # task
        if not self._wait_data_future.done():
            self._wait_data_future.set_result(None)
        self._wait_data_future = self._loop.create_future()

        return nodes, unknown_leaders_exist

    def create_builder(self):
        if self._api_version >= (0, 11):
            magic = 2
        elif self._api_version >= (0, 10):
            magic = 1
        else:
            magic = 0

        is_transactional = False
        if self._txn_manager is not None and \
                self._txn_manager.transactional_id is not None:
            is_transactional = True
        return BatchBuilder(
            magic, self._batch_size, self._compression_type,
            is_transactional=is_transactional)

    def _append_batch(self, builder, tp):
        # We must do this before actual add takes place to check for errors.
        if self._txn_manager is not None:
            self._txn_manager.maybe_add_partition_to_txn(tp)

        batch = MessageBatch(tp, builder, self._batch_ttl)
        self._batches[tp].append(batch)
        if not self._wait_data_future.done():
            self._wait_data_future.set_result(None)
        return batch

    async def add_batch(self, builder, tp, timeout):
        """Add BatchBuilder to queue by topic-partition.

        Arguments:
            builder (BatchBuilder): batch object to enqueue.
            tp (TopicPartition): topic and partition to enqueue this batch for.
            timeout (int): time in seconds to wait for a free slot in the batch
                queue.

        Returns:
            MessageBatch: delivery wrapper around the BatchBuilder object.

        Raises:
            aiokafka.errors.ProducerClosed: the accumulator has already been
                closed and flushed.
            aiokafka.errors.KafkaTimeoutError: the batch could not be added
                within the specified timeout.
        """
        if self._closed:
            raise ProducerClosed()
        if self._exception is not None:
            raise copy.copy(self._exception)

        start = time.monotonic()
        while timeout > 0:
            pending = self._batches.get(tp)
            if pending:
                await pending[-1].wait_drain(timeout=timeout)
                timeout -= time.monotonic() - start
            else:
                batch = self._append_batch(builder, tp)
                return asyncio.shield(batch.future)
        raise KafkaTimeoutError() class AIOKafkaClient:
    """Initialize an asynchronous kafka client

    Keyword Arguments:
        bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
            strings) that the consumer should contact to bootstrap initial
            cluster metadata. This does not have to be the full node list.
            It just needs to have at least one broker that will respond to
            Metadata API Request. Default port is 9092. If no servers are
            specified, will default to localhost:9092.
        client_id (str): a name for this client. This string is passed in
            each request to servers and can be used to identify specific
            server-side log entries that correspond to this client. Also
            submitted to GroupCoordinator for logging with respect to
            consumer group administration. Default: 'aiokafka-{ver}'
        request_timeout_ms (int): Client request timeout in milliseconds.
            Default: 40000.
        metadata_max_age_ms (int): The period of time in milliseconds after
            which we force a refresh of metadata even if we haven't seen
            any partition leadership changes to proactively discover any
            new brokers or partitions. Default: 300000
        retry_backoff_ms (int): Milliseconds to backoff when retrying on
            errors. Default: 100.
        api_version (str): specify which kafka API version to use.
            AIOKafka supports Kafka API versions >=0.9 only.
            If set to 'auto', will attempt to infer the broker version by
            probing various APIs. Default: auto
        security_protocol (str): Protocol used to communicate with brokers.
            Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
            Default: PLAINTEXT.
        ssl_context (ssl.SSLContext): pre-configured SSLContext for wrapping
            socket connections. For more information see :ref:`ssl_auth`.
            Default: None.
        connections_max_idle_ms (int): Close idle connections after the number
            of milliseconds specified by this config. Specifying `None` will
            disable idle checks. Default: 540000 (9 minutes).
    """

    def __init__(self, *, loop=None, bootstrap_servers='localhost',
                 client_id='aiokafka-' + __version__,
                 metadata_max_age_ms=300000,
                 request_timeout_ms=40000,
                 retry_backoff_ms=100,
                 ssl_context=None,
                 security_protocol='PLAINTEXT',
                 api_version='auto',
                 connections_max_idle_ms=540000,
                 sasl_mechanism='PLAIN',
                 sasl_plain_username=None,
                 sasl_plain_password=None,
                 sasl_kerberos_service_name='kafka',
                 sasl_kerberos_domain_name=None,
                 sasl_oauth_token_provider=None
                 ):
        if loop is None:
            loop = get_running_loop()

        if security_protocol not in (
                'SSL', 'PLAINTEXT', 'SASL_PLAINTEXT', 'SASL_SSL'):
            raise ValueError("`security_protocol` should be SSL or PLAINTEXT")
        if security_protocol in ["SSL", "SASL_SSL"] and ssl_context is None:
            raise ValueError(
                "`ssl_context` is mandatory if security_protocol=='SSL'")
        if security_protocol in ["SASL_SSL", "SASL_PLAINTEXT"]:
            if sasl_mechanism not in (
                    "PLAIN", "GSSAPI", "SCRAM-SHA-256", "SCRAM-SHA-512",
                    "OAUTHBEARER"):
                raise ValueError(
                    "only `PLAIN`, `GSSAPI`, `SCRAM-SHA-256`, "
                    "`SCRAM-SHA-512` and `OAUTHBEARER`"
                    "sasl_mechanism are supported "
                    "at the moment")
            if sasl_mechanism == "PLAIN" and \
               (sasl_plain_username is None or sasl_plain_password is None):
                raise ValueError(
                    "sasl_plain_username and sasl_plain_password required for "
                    "PLAIN sasl")

        self._bootstrap_servers = bootstrap_servers
        self._client_id = client_id
        self._metadata_max_age_ms = metadata_max_age_ms
        self._request_timeout_ms = request_timeout_ms
        if api_version != "auto":
            api_version = parse_kafka_version(api_version)
        self._api_version = api_version
        self._security_protocol = security_protocol
        self._ssl_context = ssl_context
        self._retry_backoff = retry_backoff_ms / 1000
        self._connections_max_idle_ms = connections_max_idle_ms
        self._sasl_mechanism = sasl_mechanism
        self._sasl_plain_username = sasl_plain_username
        self._sasl_plain_password = sasl_plain_password
        self._sasl_kerberos_service_name = sasl_kerberos_service_name
        self._sasl_kerberos_domain_name = sasl_kerberos_domain_name
        self._sasl_oauth_token_provider = sasl_oauth_token_provider

        self.cluster = ClusterMetadata(metadata_max_age_ms=metadata_max_age_ms)

        self._topics = set()  # empty set will fetch all topic metadata
        self._conns = {}
        self._loop = loop
        self._sync_task = None

        self._md_update_fut = None
        self._md_update_waiter = loop.create_future()
        self._get_conn_lock_value = None

    @property
    def _get_conn_lock(self):
        if self._get_conn_lock_value is None:
            self._get_conn_lock_value = asyncio.Lock()
        return self._get_conn_lock_value

    def __repr__(self):
        return '<AIOKafkaClient client_id=%s>' % self._client_id

    @property
    def api_version(self):
        if type(self._api_version) is tuple:
            return self._api_version
        # unknown api version, return minimal supported version
        return (0, 9, 0)

    @property
    def hosts(self):
        return collect_hosts(self._bootstrap_servers)

    async def close(self):
        if self._sync_task:
            self._sync_task.cancel()
            try:
                await self._sync_task
            except asyncio.CancelledError:
                pass
            self._sync_task = None
        # Be careful to wait for graceful closure of all connections, so we
        # process all pending buffers.
        futs = []
        for conn in self._conns.values():
            futs.append(conn.close(reason=CloseReason.SHUTDOWN))
        if futs:
            await asyncio.gather(*futs)

    async def bootstrap(self):
        """Try to to bootstrap initial cluster metadata"""
        assert self._loop is get_running_loop(), (
            "Please create objects with the same loop as running with"
        )
        # using request v0 for bootstrap if not sure v1 is available
        if self._api_version == "auto" or self._api_version < (0, 10):
            metadata_request = MetadataRequest[0]([])
        else:
            metadata_request = MetadataRequest[1]([])

        version_hint = None
        if self._api_version != "auto":
            version_hint = self._api_version

        for host, port, _ in self.hosts:
            log.debug("Attempting to bootstrap via node at %s:%s", host, port)

            try:
                bootstrap_conn = await create_conn(
                    host, port, client_id=self._client_id,
                    request_timeout_ms=self._request_timeout_ms,
                    ssl_context=self._ssl_context,
                    security_protocol=self._security_protocol,
                    max_idle_ms=self._connections_max_idle_ms,
                    sasl_mechanism=self._sasl_mechanism,
                    sasl_plain_username=self._sasl_plain_username,
                    sasl_plain_password=self._sasl_plain_password,
                    sasl_kerberos_service_name=self._sasl_kerberos_service_name,  # noqa: ignore=E501
                    sasl_kerberos_domain_name=self._sasl_kerberos_domain_name,
                    sasl_oauth_token_provider=self._sasl_oauth_token_provider,
                    version_hint=version_hint)
            except (OSError, asyncio.TimeoutError) as err:
                log.error('Unable connect to "%s:%s": %s', host, port, err)
                continue

            try:
                metadata = await bootstrap_conn.send(metadata_request)
            except (KafkaError, asyncio.TimeoutError) as err:
                log.warning('Unable to request metadata from "%s:%s": %s',
                            host, port, err)
                bootstrap_conn.close()
                continue

            self.cluster.update_metadata(metadata)

            # A cluster with no topics can return no broker metadata...
            # In that case, we should keep the bootstrap connection till
            # we get a normal cluster layout.
            if not len(self.cluster.brokers()):
                bootstrap_id = ('bootstrap', ConnectionGroup.DEFAULT)
                self._conns[bootstrap_id] = bootstrap_conn
            else:
                bootstrap_conn.close()

            log.debug('Received cluster metadata: %s', self.cluster)
            break
        else:
            raise KafkaConnectionError(
                f'Unable to bootstrap from {self.hosts}')

        # detect api version if need
        if self._api_version == 'auto':
            self._api_version = await self.check_version()

        if self._sync_task is None:
            # starting metadata synchronizer task
            self._sync_task = create_task(self._md_synchronizer())

    async def _md_synchronizer(self):
        """routine (async task) for synchronize cluster metadata every
        `metadata_max_age_ms` milliseconds"""
        while True:
            await asyncio.wait(
                [self._md_update_waiter],
                timeout=self._metadata_max_age_ms / 1000)

            topics = self._topics
            if self._md_update_fut is None:
                self._md_update_fut = create_future()
            ret = await self._metadata_update(self.cluster, topics)
            # If list of topics changed during metadata update we must update
            # it again right away.
            if topics != self._topics:
                continue
            # Earlier this waiter was set before sending metadata_request,
            # but that was to avoid topic list changes being unnoticed, which
            # is handled explicitly now.
            self._md_update_waiter = create_future()

            self._md_update_fut.set_result(ret)
            self._md_update_fut = None

    def get_random_node(self):
        """choice random node from known cluster brokers

        Returns:
            nodeId - identifier of broker
        """
        nodeids = [b.nodeId for b in self.cluster.brokers()]
        if not nodeids:
            return None
        return random.choice(nodeids)

    async def _metadata_update(self, cluster_metadata, topics):
        assert isinstance(cluster_metadata, ClusterMetadata)
        topics = list(topics)
        version_id = 0 if self.api_version < (0, 10) else 1
        if version_id == 1 and not topics:
            topics = None
        metadata_request = MetadataRequest[version_id](topics)
        nodeids = [b.nodeId for b in self.cluster.brokers()]
        bootstrap_id = ('bootstrap', ConnectionGroup.DEFAULT)
        if bootstrap_id in self._conns:
            nodeids.append('bootstrap')
        random.shuffle(nodeids)
        for node_id in nodeids:
            conn = await self._get_conn(node_id)

            if conn is None:
                continue
            log.debug("Sending metadata request %s to node %s",
                      metadata_request, node_id)

            try:
                metadata = await conn.send(metadata_request)
            except (KafkaError, asyncio.TimeoutError) as err:
                log.error(
                    'Unable to request metadata from node with id %s: %r',
                    node_id, err)
                continue

            # don't update the cluster if there are no valid nodes...the topic
            # we want may still be in the process of being created which means
            # we will get errors and no nodes until it exists
            if not metadata.brokers:
                return False

            cluster_metadata.update_metadata(metadata)

            # We only keep bootstrap connection to update metadata until
            # proper cluster layout is available.
            if bootstrap_id in self._conns and len(self.cluster.brokers()):
                conn = self._conns.pop(bootstrap_id)
                conn.close()

            break
        else:
            log.error('Unable to update metadata from %s', nodeids)
            cluster_metadata.failed_update(None)
            return False
        return True

    def force_metadata_update(self):
        """Update cluster metadata

        Returns:
            True/False - metadata updated or not
        """
        if self._md_update_fut is None:
            # Wake up the `_md_synchronizer` task
            if not self._md_update_waiter.done():
                self._md_update_waiter.set_result(None)
            self._md_update_fut = self._loop.create_future()
        # Metadata will be updated in the background by syncronizer
        return asyncio.shield(self._md_update_fut)

    async def fetch_all_metadata(self):
        cluster_md = ClusterMetadata(
            metadata_max_age_ms=self._metadata_max_age_ms)
        updated = await self._metadata_update(cluster_md, [])
        if not updated:
            raise KafkaError(
                'Unable to get cluster metadata over all known brokers')
        return cluster_md

    def add_topic(self, topic):
        """Add a topic to the list of topics tracked via metadata.

        Arguments:
            topic (str): topic to track
        """
        if topic in self._topics:
            res = self._loop.create_future()
            res.set_result(True)
        else:
            res = self.force_metadata_update()
        self._topics.add(topic)
        return res

    def set_topics(self, topics):
        """Set specific topics to track for metadata.

        Arguments:
            topics (list of str): topics to track
        """
        assert not isinstance(topics, str)
        if not topics or set(topics).difference(self._topics):
            res = self.force_metadata_update()
        else:
            res = self._loop.create_future()
            res.set_result(True)
        self._topics = set(topics)
        return res

    def _on_connection_closed(self, conn, reason):
        """ Callback called when connection is closed
        """
        # Connection failures imply that our metadata is stale, so let's
        # refresh
        if reason == CloseReason.CONNECTION_BROKEN or \
                reason == CloseReason.CONNECTION_TIMEOUT:
            self.force_metadata_update()

    async def _get_conn(
        self, node_id, *, group=ConnectionGroup.DEFAULT,
        no_hint=False
    ):
        "Get or create a connection to a broker using host and port"
        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

        try:
            if group == ConnectionGroup.DEFAULT:
                broker = self.cluster.broker_metadata(node_id)
                # XXX: earlier we only did an assert here, but it seems it's
                # possible to get a leader that is for some reason not in
                # metadata.
                # I think requerying metadata should solve this problem
                if broker is None:
                    raise StaleMetadata(
                        'Broker id %s not in current metadata' % node_id)
            else:
                broker = self.cluster.coordinator_metadata(node_id)
                assert broker is not None

            log.debug("Initiating connection to node %s at %s:%s",
                      node_id, broker.host, broker.port)

            async with self._get_conn_lock:
                if conn_id in self._conns:
                    return self._conns[conn_id]

                version_hint = self._api_version
                if version_hint == "auto" or no_hint:
                    version_hint = None

                self._conns[conn_id] = await create_conn(
                    broker.host, broker.port,
                    client_id=self._client_id,
                    request_timeout_ms=self._request_timeout_ms,
                    ssl_context=self._ssl_context,
                    security_protocol=self._security_protocol,
                    on_close=self._on_connection_closed,
                    max_idle_ms=self._connections_max_idle_ms,
                    sasl_mechanism=self._sasl_mechanism,
                    sasl_plain_username=self._sasl_plain_username,
                    sasl_plain_password=self._sasl_plain_password,
                    sasl_kerberos_service_name=self._sasl_kerberos_service_name,  # noqa: ignore=E501
                    sasl_kerberos_domain_name=self._sasl_kerberos_domain_name,
                    sasl_oauth_token_provider=self._sasl_oauth_token_provider,
                    version_hint=version_hint
                )
        except (OSError, asyncio.TimeoutError, KafkaError) as err:
            log.error('Unable connect to node with id %s: %s', node_id, err)
            if group == ConnectionGroup.DEFAULT:
                # Connection failures imply that our metadata is stale, so
                # let's refresh
                self.force_metadata_update()
            return None
        else:
            return self._conns[conn_id]

    async def ready(self, node_id, *, group=ConnectionGroup.DEFAULT):
        conn = await self._get_conn(node_id, group=group)
        if conn is None:
            return False
        return True

    async def send(self, node_id, request, *, group=ConnectionGroup.DEFAULT):
        """Send a request to a specific node.

        Arguments:
            node_id (int): destination node
            request (Struct): request object (not-encoded)

        Raises:
            kafka.errors.RequestTimedOutError
            kafka.errors.NodeNotReadyError
            kafka.errors.KafkaConnectionError
            kafka.errors.CorrelationIdError

        Returns:
            Future: resolves to Response struct
        """
        if not (await self.ready(node_id, group=group)):
            raise NodeNotReadyError(
                "Attempt to send a request to node"
                " which is not ready (node id {}).".format(node_id))

        # Every request gets a response, except one special case:
        expect_response = True
        if isinstance(request, tuple(ProduceRequest)) and \
                request.required_acks == 0:
            expect_response = False

        future = self._conns[(node_id, group)].send(
            request, expect_response=expect_response)
        try:
            result = await future
        except asyncio.TimeoutError:
            # close connection so it is renewed in next request
            self._conns[(node_id, group)].close(
                reason=CloseReason.CONNECTION_TIMEOUT)
            raise RequestTimedOutError()
        else:
            return result

    async def check_version(self, node_id=None):
        """Attempt to guess the broker version"""
        if node_id is None:
            default_group_conns = [
                n_id for (n_id, group) in self._conns.keys()
                if group == ConnectionGroup.DEFAULT
            ]
            if default_group_conns:
                node_id = default_group_conns[0]
            else:
                assert self.cluster.brokers(), 'no brokers in metadata'
                node_id = list(self.cluster.brokers())[0].nodeId

        from kafka.protocol.admin import (
            ListGroupsRequest_v0, ApiVersionRequest_v0)
        from kafka.protocol.commit import (
            OffsetFetchRequest_v0, GroupCoordinatorRequest_v0)
        from kafka.protocol.metadata import MetadataRequest_v0
        test_cases = [
            ((0, 10), ApiVersionRequest_v0()),
            ((0, 9), ListGroupsRequest_v0()),
            ((0, 8, 2), GroupCoordinatorRequest_v0('aiokafka-default-group')),
            ((0, 8, 1), OffsetFetchRequest_v0('aiokafka-default-group', [])),
            ((0, 8, 0), MetadataRequest_v0([])),
        ]

        # kafka kills the connection when it does not recognize an API request
        # so we can send a test request and then follow immediately with a
        # vanilla MetadataRequest. If the server did not recognize the first
        # request, both will be failed with a ConnectionError that wraps
        # socket.error (32, 54, or 104)
        conn = await self._get_conn(node_id, no_hint=True)
        if conn is None:
            raise KafkaConnectionError(
                f"No connection to node with id {node_id}")
        for version, request in test_cases:
            try:
                if not conn.connected():
                    await conn.connect()
                assert conn, f'no connection to node with id {node_id}'
                # request can be ignored by Kafka broker,
                # so we send metadata request and wait response
                task = create_task(conn.send(request))
                await asyncio.wait([task], timeout=0.1)
                try:
                    await conn.send(MetadataRequest_v0([]))
                except KafkaError:
                    # metadata request can be cancelled in case
                    # of invalid correlationIds order
                    pass
                response = await task
            except KafkaError:
                continue
            else:
                # To avoid having a connection in undefined state
                if node_id != "bootstrap" and conn.connected():
                    conn.close()
                if isinstance(request, ApiVersionRequest_v0):
                    # Starting from 0.10 kafka broker we determine version
                    # by looking at ApiVersionResponse
                    version = self._check_api_version_response(response)
                return version

        raise UnrecognizedBrokerVersion()

    def _check_api_version_response(self, response):
        # The logic here is to check the list of supported request versions
        # in descending order. As soon as we find one that works, return it
        test_cases = [
            # format (<broker version>, <needed struct>)
            # TODO Requires unreleased version of python-kafka
            # ((2, 6, 0), DescribeClientQuotasRequest[0]),
            ((2, 5, 0), DescribeAclsRequest_v2),
            ((2, 4, 0), ProduceRequest[8]),
            ((2, 3, 0), FetchRequest[11]),
            ((2, 2, 0), OffsetRequest[5]),
            ((2, 1, 0), FetchRequest[10]),
            ((2, 0, 0), FetchRequest[8]),
            ((1, 1, 0), FetchRequest[7]),
            ((1, 0, 0), MetadataRequest[5]),
            ((0, 11, 0), MetadataRequest[4]),
            ((0, 10, 2), OffsetFetchRequest[2]),
            ((0, 10, 1), MetadataRequest[2]),
        ]

        error_type = Errors.for_code(response.error_code)
        assert error_type is Errors.NoError, "API version check failed"
        max_versions = {
            api_key: max_version
            for api_key, _, max_version in response.api_versions
        }
        # Get the best match of test cases
        for broker_version, struct in test_cases:
            if max_versions.get(struct.API_KEY, -1) >= struct.API_VERSION:
                return broker_version

        # We know that ApiVersionResponse is only supported in 0.10+
        # so if all else fails, choose that
        return (0, 10, 0)

    async def _wait_on_metadata(self, topic):
        """
        Wait for cluster metadata including partitions for the given topic to
        be available.

        Arguments:
            topic (str): topic we want metadata for

        Returns:
            set: partition ids for the topic

        Raises:
            UnknownTopicOrPartitionError: if no topic or partitions found
                in cluster metadata
        """
        if topic in self.cluster.topics():
            return self.cluster.partitions_for_topic(topic)

        # add topic to metadata topic list if it is not there already.
        self.add_topic(topic)

        t0 = time.monotonic()
        while True:
            await self.force_metadata_update()
            if topic in self.cluster.topics():
                break
            if (time.monotonic() - t0) > (self._request_timeout_ms / 1000):
                raise UnknownTopicOrPartitionError()
            if topic in self.cluster.unauthorized_topics:
                raise Errors.TopicAuthorizationFailedError(topic)
            await asyncio.sleep(self._retry_backoff)

        return self.cluster.partitions_for_topic(topic)

    async def _maybe_wait_metadata(self):
        if self._md_update_fut is not None:
            await asyncio.shield(self._md_update_fut)

    async def coordinator_lookup(self, coordinator_type, coordinator_key):
        """ Lookup which node in the cluster is the coordinator for a certain
        role (Transaction coordinator or Group coordinator atm.)
        NOTE: Client keeps track of all coordination nodes separately, as they
        all have different sockets and ids.
        """

        node_id = self.get_random_node()
        assert node_id is not None, "Did we not perform bootstrap?"

        log.debug(
            "Sending FindCoordinator request for key %s to broker %s",
            coordinator_key, node_id)

        if self.api_version > (0, 11):
            request = FindCoordinatorRequest[1](
                coordinator_key, coordinator_type)
        else:
            # Group coordination only
            assert coordinator_type == CoordinationType.GROUP, \
                "No transactions for older brokers"
            request = FindCoordinatorRequest[0](coordinator_key)

        resp = await self.send(node_id, request)
        log.debug("Received group coordinator response %s", resp)
        error_type = Errors.for_code(resp.error_code)
        if error_type is not Errors.NoError:
            err = error_type()
            raise err
        self.cluster.add_coordinator(
            resp.coordinator_id, resp.host, resp.port, rack=None,
            purpose=(coordinator_type, coordinator_key))
        return resp.coordinator_id class CoordinationType:

    GROUP = 0
    TRANSACTION = 1 class ConnectionGroup:

    DEFAULT = 0
    COORDINATION = 1 class RecordMetadata(NamedTuple):
class ConsumerRecord(Generic[KT, VT]):
class OffsetAndTimestamp(NamedTuple):
KT = TypeVar("KT")
VT = TypeVar("VT") def get_running_loop() -> asyncio.AbstractEventLoop:
    loop = asyncio.get_event_loop()
    if not loop.is_running():
        raise RuntimeError(
            "The object should be created within an async function or "
            "provide loop directly."
        )
    return loop class CoordinatorNotAvailableError(GroupCoordinatorNotAvailableError):
class NotCoordinatorError(NotCoordinatorForGroupError):
class CoordinatorLoadInProgressError(GroupLoadInProgressError):
class ConsumerStoppedError(Exception):
class IllegalOperation(Exception):
class NoOffsetForPartitionError(KafkaError):
class RecordTooLargeError(KafkaError):
class ProducerClosed(KafkaError):
class ProducerFenced(KafkaError):
class OutOfOrderSequenceNumber(BrokerResponseError):
class DuplicateSequenceNumber(BrokerResponseError):
class InvalidProducerEpoch(BrokerResponseError):
class InvalidTxnState(BrokerResponseError):
class InvalidProducerIdMapping(BrokerResponseError):
class InvalidTransactionTimeout(BrokerResponseError):
class ConcurrentTransactions(BrokerResponseError):
class TransactionCoordinatorFenced(BrokerResponseError):
class TransactionalIdAuthorizationFailed(BrokerResponseError):
class SecurityDisabled(BrokerResponseError):
class OperationNotAttempted(BrokerResponseError):
class KafkaStorageError(BrokerResponseError):
class LogDirNotFound(BrokerResponseError):
class SaslAuthenticationFailed(BrokerResponseError):
class UnknownProducerId(BrokerResponseError):
class ReassignmentInProgress(BrokerResponseError):
class DelegationTokenAuthDisabled(BrokerResponseError):
class DelegationTokenNotFound(BrokerResponseError):
class DelegationTokenOwnerMismatch(BrokerResponseError):
class DelegationTokenRequestNotAllowed(BrokerResponseError):
class DelegationTokenAuthorizationFailed(BrokerResponseError):
class DelegationTokenExpired(BrokerResponseError):
class InvalidPrincipalType(BrokerResponseError):
class NonEmptyGroup(BrokerResponseError):
class GroupIdNotFound(BrokerResponseError):
class FetchSessionIdNotFound(BrokerResponseError):
class InvalidFetchSessionEpoch(BrokerResponseError):
class ListenerNotFound(BrokerResponseError):
    def __init__(
        self,
        msg="There is a newer producer using the same transactional_id or"
            "transaction timeout occurred (check that processing time is "
            "below transaction_timeout_ms)"
    ):
def _iter_broker_errors():
def for_code(error_code):
from unittest import mock
from ._testutil import (
    KafkaIntegrationTestCase, run_until_complete, kafka_versions
)
from kafka.protocol.produce import ProduceRequest, ProduceResponse
from aiokafka.producer.sender import (
    Sender, InitPIDHandler, AddPartitionsToTxnHandler,
    AddOffsetsToTxnHandler, TxnOffsetCommitHandler, EndTxnHandler,
    BaseHandler, SendProduceReqHandler
)
from aiokafka.producer.transaction_manager import (
    TransactionManager, TransactionState
)
from aiokafka.protocol.transaction import (
    InitProducerIdRequest, InitProducerIdResponse,
    AddPartitionsToTxnRequest, AddPartitionsToTxnResponse,
    AddOffsetsToTxnRequest, AddOffsetsToTxnResponse,
    TxnOffsetCommitRequest, TxnOffsetCommitResponse,
    EndTxnRequest, EndTxnResponse
)
from aiokafka.producer.message_accumulator import MessageAccumulator
from aiokafka.client import AIOKafkaClient, CoordinationType, ConnectionGroup
from aiokafka.structs import TopicPartition, OffsetAndMetadata
from aiokafka.util import get_running_loop
from aiokafka.errors import (
    NoError, UnknownError,
    CoordinatorNotAvailableError, NotCoordinatorError,
    CoordinatorLoadInProgressError, ConcurrentTransactions,
    UnknownTopicOrPartitionError, InvalidProducerEpoch,
    ProducerFenced, InvalidProducerIdMapping, InvalidTxnState,
    RequestTimedOutError, DuplicateSequenceNumber, KafkaError,
    TopicAuthorizationFailedError, OperationNotAttempted,
    TransactionalIdAuthorizationFailed, GroupAuthorizationFailedError
)
from kafka.protocol.metadata import MetadataRequest
        tm.offset_committed.assert_not_called()
        self.assertEqual(cm.exception.args[0], "test_tid")

        # Handle unknown error
        resp = create_response(UnknownError)
        with self.assertRaises(UnknownError):
            add_handler.handle_response(resp)
        tm.offset_committed.assert_not_called()

        # GroupAuthorizationFailedError case
        self.assertNotEqual(tm.state, TransactionState.ABORTABLE_ERROR)
        resp = create_response(GroupAuthorizationFailedError)
        backoff = add_handler.handle_response(resp)
        self.assertIsNone(backoff)
        tm.offset_committed.assert_not_called()
        with self.assertRaises(GroupAuthorizationFailedError) as cm:
            tm.committing_transaction()
        self.assertEqual(cm.exception.args[0], "some_group")
        self.assertEqual(tm.state, TransactionState.ABORTABLE_ERROR)

    @run_until_complete
    async def test_sender__do_end_txn_create(self):
        sender = await self._setup_sender()
        add_handler = EndTxnHandler(sender, 0)

        req = add_handler.create_request()
        self.assertEqual(req.API_KEY, EndTxnRequest[0].API_KEY)
        self.assertEqual(req.API_VERSION, 0)
        self.assertEqual(req.transactional_id, "test_tid")
        self.assertEqual(req.producer_id, 120)
        self.assertEqual(req.producer_epoch, 22)
        self.assertEqual(req.transaction_result, 0)

    @run_until_complete
    async def test_sender__do_end_txn_ok(self):
        sender = await self._setup_sender()
        add_handler = EndTxnHandler(sender, 0)
        tm = sender._txn_manager
        tm.complete_transaction = mock.Mock()

        # Handle response
        cls = EndTxnResponse[0]
        resp = cls(
            throttle_time_ms=300,
            error_code=NoError.errno
        )
        backoff = add_handler.handle_response(resp)
        self.assertIsNone(backoff)
        tm.complete_transaction.assert_called_with()

    @run_until_complete
    async def test_sender__do_end_txn_not_ok(self):
        sender = await self._setup_sender()
        add_handler = EndTxnHandler(sender, 0)
        tm = sender._txn_manager
        tm.complete_transaction = mock.Mock()

        def create_response(error_type):
            cls = EndTxnResponse[0]
            resp = cls(
                throttle_time_ms=300,
                error_code=error_type.errno
            )
            return resp

        # Handle coordination errors
        for error_cls in [CoordinatorNotAvailableError, NotCoordinatorError]:
            with mock.patch.object(sender, "_coordinator_dead") as mocked:
                resp = create_response(error_cls)
                backoff = add_handler.handle_response(resp)
                self.assertEqual(backoff, 0.1)
                tm.complete_transaction.assert_not_called()
                mocked.assert_called_with(CoordinationType.TRANSACTION)

        # Not coordination retriable errors
        for error_cls in [
                CoordinatorLoadInProgressError, ConcurrentTransactions]:
            resp = create_response(error_cls)
            backoff = add_handler.handle_response(resp)
            self.assertEqual(backoff, 0.1)
            tm.complete_transaction.assert_not_called()

        # ProducerFenced case
        resp = create_response(InvalidProducerEpoch)
        with self.assertRaises(ProducerFenced):
            add_handler.handle_response(resp)
        tm.complete_transaction.assert_not_called()

        for error_type in [InvalidTxnState]:
            resp = create_response(error_type)
            with self.assertRaises(error_type):
                add_handler.handle_response(resp)
            tm.complete_transaction.assert_not_called()

        # Handle unknown error
        resp = create_response(UnknownError)
        with self.assertRaises(UnknownError):
            add_handler.handle_response(resp)
        tm.complete_transaction.assert_not_called()

    @run_until_complete
    async def test_sender__produce_request_create(self):
        sender = await self._setup_sender()
        tp = TopicPartition("my_topic", 0)
        batch_mock = mock.Mock()
        batch_mock.get_data_buffer = mock.Mock(return_value=b"123")
        send_handler = SendProduceReqHandler(
            sender, {tp: batch_mock})

        req = send_handler.create_request()
        self.assertEqual(req.API_KEY, ProduceRequest[0].API_KEY)
        if req.API_VERSION > 2:
            self.assertEqual(req.transactional_id, "test_tid")
        self.assertEqual(req.required_acks, -1)
        self.assertEqual(req.timeout, 40000)
        self.assertEqual(req.topics, [("my_topic", [(0, b"123")])])

    @run_until_complete
    async def test_sender__produce_request_ok(self):
        sender = await self._setup_sender()

        tp = TopicPartition("my_topic", 0)