Repository: apache/kafka
File: tests/kafkatest/services/verifiable_consumer.py (Python)

Code before:
        assignment = []
        for topic_partition in event["partitions"]:
            tp = _create_partition_from_dict(topic_partition)
            assignment.append(tp)
        logger.debug("Partitions %s assigned to %s" % (assignment, node.account.hostname))
        self.assignment.extend(assignment)

# This needs to be used for consumer protocol.
class ConsumerProtocolConsumerEventHandler(IncrementalAssignmentConsumerEventHandler):
    def __init__(self, node, verify_offsets, idx):
        super().__init__(node, verify_offsets, idx)

    def handle_partitions_revoked(self, event, node, logger):
        self.revoked_count += 1
        self.position = {}
        revoked = []

        for topic_partition in event["partitions"]:
            tp = _create_partition_from_dict(topic_partition)
            # tp existing in self.assignment is not guaranteed in the new consumer
            # if it shuts down when revoking partitions for reconciliation.
            if tp in self.assignment:
                self.assignment.remove(tp)
            revoked.append(tp)

        logger.debug("Partitions %s revoked from %s" % (revoked, node.account.hostname))

class VerifiableConsumer(KafkaPathResolverMixin, VerifiableClientMixin, BackgroundThreadService):
    """This service wraps org.apache.kafka.tools.VerifiableConsumer for use in
    system testing. 
    
    NOTE: this class should be treated as a PUBLIC API. Downstream users use
    this service both directly and through class extension, so care must be 
    taken to ensure compatibility.
    """

    PERSISTENT_ROOT = "/mnt/verifiable_consumer"
    STDOUT_CAPTURE = os.path.join(PERSISTENT_ROOT, "verifiable_consumer.stdout")
    STDERR_CAPTURE = os.path.join(PERSISTENT_ROOT, "verifiable_consumer.stderr")
    LOG_DIR = os.path.join(PERSISTENT_ROOT, "logs")
    LOG_FILE = os.path.join(LOG_DIR, "verifiable_consumer.log")
    LOG4J_CONFIG = os.path.join(PERSISTENT_ROOT, "tools-log4j.properties")
    CONFIG_FILE = os.path.join(PERSISTENT_ROOT, "verifiable_consumer.properties")

    logs = {
        "verifiable_consumer_stdout": {
            "path": STDOUT_CAPTURE,
            "collect_default": False},
        "verifiable_consumer_stderr": {
            "path": STDERR_CAPTURE,
            "collect_default": False},

Review comment (bug):
I am not sure to follow this comment. A partition cannot not be revoked if it was not assigned to start with. Have you see such failures? Could you elaborate a bit more on it please?

Code after:
        self.assigned_count += 1
        self.state = ConsumerState.Joined
        assignment = []
        for topic_partition in event["partitions"]:
            tp = _create_partition_from_dict(topic_partition)
            assignment.append(tp)
        logger.debug("Partitions %s assigned to %s" % (assignment, node.account.hostname))
        self.assignment.extend(assignment)

# This needs to be used for consumer protocol.
class ConsumerProtocolConsumerEventHandler(IncrementalAssignmentConsumerEventHandler):
    def __init__(self, node, verify_offsets, idx, **kwargs):
        super().__init__(node, verify_offsets, idx, **kwargs)

    def handle_partitions_revoked(self, event, node, logger):
        # The handler state is not transitioned to Rebalancing as the records can only be
        # consumed in Joined state (see ConsumerEventHandler.handle_records_consumed).
        # The consumer with consumer protocol should still be able to consume messages during rebalance.
        self.revoked_count += 1
        self.position = {}
        revoked = []

        for topic_partition in event["partitions"]:
            tp = _create_partition_from_dict(topic_partition)
            # tp existing in self.assignment is not guaranteed in the new consumer
            # if it shuts down when revoking partitions for reconciliation.
            if tp in self.assignment:
                self.assignment.remove(tp)
            revoked.append(tp)

        logger.debug("Partitions %s revoked from %s" % (revoked, node.account.hostname))

class VerifiableConsumer(KafkaPathResolverMixin, VerifiableClientMixin, BackgroundThreadService):
    """This service wraps org.apache.kafka.tools.VerifiableConsumer for use in
    system testing. 
    
    NOTE: this class should be treated as a PUBLIC API. Downstream users use
    this service both directly and through class extension, so care must be 
    taken to ensure compatibility.
    """

    PERSISTENT_ROOT = "/mnt/verifiable_consumer"
    STDOUT_CAPTURE = os.path.join(PERSISTENT_ROOT, "verifiable_consumer.stdout")
    STDERR_CAPTURE = os.path.join(PERSISTENT_ROOT, "verifiable_consumer.stderr")
    LOG_DIR = os.path.join(PERSISTENT_ROOT, "logs")
    LOG_FILE = os.path.join(LOG_DIR, "verifiable_consumer.log")
    LOG4J_CONFIG = os.path.join(PERSISTENT_ROOT, "tools-log4j.properties")
    CONFIG_FILE = os.path.join(PERSISTENT_ROOT, "verifiable_consumer.properties")

    logs = {
        "verifiable_consumer_stdout": {