Repository: apache/kafka
File: tests/kafkatest/tests/client/consumer_test.py (Python)

Code before:
    )
    @matrix(
        clean_shutdown=[True],
        enable_autocommit=[True, False],
        metadata_quorum=[quorum.isolated_kraft],
        use_new_coordinator=[True],
        group_protocol=consumer_group.all_group_protocols
    )
    def test_consumer_failure(self, clean_shutdown, enable_autocommit, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
        partition = TopicPartition(self.TOPIC, 0)

        consumer = self.setup_consumer(self.TOPIC, enable_autocommit=enable_autocommit, group_protocol=group_protocol)
        producer = self.setup_producer(self.TOPIC)

        consumer.start()
        self.await_all_members(consumer)

        partition_owner_container = {partition: None}
        def check_partition_owner(partition_owner_container):
            partition_owner_container[partition] = consumer.owner(partition)
            return partition_owner_container[partition] is not None

        wait_until(lambda: check_partition_owner(partition_owner_container),
                   timeout_sec=self.session_timeout_sec*2+5,
                   err_msg="Timed out waiting for partition to be assigned.")
        partition_owner = partition_owner_container[partition]

        # startup the producer and ensure that some records have been written
        producer.start()
        self.await_produced_messages(producer)

        # stop the partition owner and await its shutdown
        consumer.kill_node(partition_owner, clean_shutdown=clean_shutdown)
        wait_until(lambda: len(consumer.joined_nodes()) == (self.num_consumers - 1) and consumer.owner(partition) is not None,
                   timeout_sec=self.session_timeout_sec*2+5,
                   err_msg="Timed out waiting for consumer to close")

        # ensure that the remaining consumer does some work after rebalancing
        self.await_consumed_messages(consumer, min_messages=1000)

        consumer.stop_all()

        if clean_shutdown:
            # if the total records consumed matches the current position, we haven't seen any duplicates
            # this can only be guaranteed with a clean shutdown
            assert consumer.current_position(partition) == consumer.total_consumed(), \
                "Total consumed records %d did not match consumed position %d" % \
                (consumer.total_consumed(), consumer.current_position(partition))
        else:
            # we may have duplicates in a hard failure
            assert consumer.current_position(partition) <= consumer.total_consumed(), \

Review comment (suggestion):
For new protocol consumer, it's possible that all members have joined but the members are not fully stablized and the partition currently doesn't have an owner. We need to wait for a while for the assertion.

It's a bit complicated because the partition owner will be used on L433 so we want to keep a variable here.

Code after:
        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
        use_new_coordinator=[False]
    )
    @matrix(
        clean_shutdown=[True],
        enable_autocommit=[True, False],
        metadata_quorum=[quorum.isolated_kraft],
        use_new_coordinator=[True],
        group_protocol=consumer_group.all_group_protocols
    )
    def test_consumer_failure(self, clean_shutdown, enable_autocommit, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
        partition = TopicPartition(self.TOPIC, 0)

        consumer = self.setup_consumer(self.TOPIC, enable_autocommit=enable_autocommit, group_protocol=group_protocol)
        producer = self.setup_producer(self.TOPIC)

        consumer.start()
        self.await_all_members(consumer)
        self.await_all_members_stabilized(self.TOPIC, self.NUM_PARTITIONS, consumer, timeout_sec=60)
        partition_owner = consumer.owner(partition)

        # startup the producer and ensure that some records have been written
        producer.start()
        self.await_produced_messages(producer)

        # stop the partition owner and await its shutdown
        consumer.kill_node(partition_owner, clean_shutdown=clean_shutdown)
        wait_until(lambda: len(consumer.joined_nodes()) == (self.num_consumers - 1) and consumer.owner(partition) is not None,
                   timeout_sec=self.session_timeout_sec*2+5,
                   err_msg="Timed out waiting for consumer to close")

        # ensure that the remaining consumer does some work after rebalancing
        self.await_consumed_messages(consumer, min_messages=1000)

        consumer.stop_all()

        if clean_shutdown:
            # if the total records consumed matches the current position, we haven't seen any duplicates
            # this can only be guaranteed with a clean shutdown
            assert consumer.current_position(partition) == consumer.total_consumed(), \
                "Total consumed records %d did not match consumed position %d" % \
                (consumer.total_consumed(), consumer.current_position(partition))
        else:
            # we may have duplicates in a hard failure
            assert consumer.current_position(partition) <= consumer.total_consumed(), \
                "Current position %d greater than the total number of consumed records %d" % \
                (consumer.current_position(partition), consumer.total_consumed())

        # if autocommit is not turned on, we can also verify the last committed offset
        if not enable_autocommit:
            assert consumer.last_commit(partition) == consumer.current_position(partition), \