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

Code before:
            assert tp in self.assignment, \
                "Topic partition %s cannot be revoked from %s as it was not previously assigned to that consumer" % \
                (tp, node.account.hostname)
            self.assignment.remove(tp)
            revoked.append(tp)

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

    def handle_partitions_assigned(self, event, node, logger):
        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):
        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")

Review comment (suggestion):
I wonder if we should add a comment explaining why we don't transition to `ConsumerState.Rebalancing` in this event handler.

Code after:
        for topic_partition in event["partitions"]:
            tp = _create_partition_from_dict(topic_partition)
            assert tp in self.assignment, \
                "Topic partition %s cannot be revoked from %s as it was not previously assigned to that consumer" % \
                (tp, node.account.hostname)
            self.assignment.remove(tp)
            revoked.append(tp)

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

    def handle_partitions_assigned(self, event, node, logger):
        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.
    """