Package com.linkedin.venice.unit.kafka.consumer

Class MockInMemoryConsumer

java.lang.Object
com.linkedin.venice.unit.kafka.consumer.MockInMemoryConsumer
All Implemented Interfaces:
PubSubConsumerAdapter, Closeable, AutoCloseable
public class MockInMemoryConsumer extends Object implements PubSubConsumerAdapter
A PubSubConsumerAdapter implementation which reads messages from the InMemoryKafkaBroker. Used in unit tests as a lightweight alternative to a full-fledged integration test. Can be configured with various PollStrategy implementations in order to tweak the consuming behavior. When MockInMemoryConsumer is used to simulate the shared consumer behavior, there might be 2 different threads calling the methods from this class. For example, consumer task thread from KafkaConsumerService will periodically call poll(long) and StoreIngestionTask thread is calling resetOffset(PubSubTopicPartition), which may cause test failure. TODO: Remove synchronized keyword in this class when consumer operations in consumption task is event-driven.

  • Constructor Details

    • MockInMemoryConsumer

      public MockInMemoryConsumer(InMemoryKafkaBroker broker, PollStrategy pollStrategy, PubSubConsumerAdapter delegate)
      Parameters:
      delegate - Can be used to pass a mock, in order to verify calls. Note: functions that return do not return the result of the mock, but rather the results of the in-memory consumer components.

  • Method Details

    • subscribe

      public void subscribe(PubSubTopicPartition pubSubTopicPartition, long lastReadOffset)
      Description copied from interface: PubSubConsumerAdapter
      Subscribes to a topic-partition if it is not already subscribed. If the topic-partition is already subscribed, this method is a no-op. The method assumes that the topic-partition exists.
      Specified by:
      subscribe in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The topic-partition to subscribe to.
      lastReadOffset - The last read offset for the topic-partition. A poll call following a subscribe call will return messages from the offset (lastReadOffset + 1).

    • unSubscribe

      public void unSubscribe(PubSubTopicPartition pubSubTopicPartition)
      Description copied from interface: PubSubConsumerAdapter
      Unsubscribes the consumer from a specified topic-partition. If the consumer was previously subscribed to the given partition, it will be unsubscribed, and the associated partition assignments and tracked offsets will be updated accordingly.
      Specified by:
      unSubscribe in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic-partition to unsubscribe from.

    • batchUnsubscribe

      public void batchUnsubscribe(Set<PubSubTopicPartition> pubSubTopicPartitionSet)
      Description copied from interface: PubSubConsumerAdapter
      Unsubscribes the consumer from a batch of topic-partitions.
      Specified by:
      batchUnsubscribe in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartitionSet - A set of topic-partitions to unsubscribe from.

    • resetOffset

      public void resetOffset(PubSubTopicPartition pubSubTopicPartition)
      Description copied from interface: PubSubConsumerAdapter
      Resets the offset for a specific topic-partition, seeking it to the beginning of the topic. If the topic partition is not currently subscribed, a PubSubUnsubscribedTopicPartitionException is thrown.
      Specified by:
      resetOffset in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic-partition for which to reset the offset.

    • close

      public void close()
      Description copied from interface: PubSubConsumerAdapter
      Closes the PubSub consumer and releases any associated resources.
      Specified by:
      close in interface AutoCloseable
      Specified by:
      close in interface Closeable
      Specified by:
      close in interface PubSubConsumerAdapter

    • poll

      public Map<PubSubTopicPartition,List<PubSubMessage<KafkaKey,KafkaMessageEnvelope,Long>>> poll(long timeout)
      Description copied from interface: PubSubConsumerAdapter
      Polls the Kafka consumer for messages from the subscribed topic partitions within the specified time duration.
      Specified by:
      poll in interface PubSubConsumerAdapter
      Parameters:
      timeout - The maximum time, in milliseconds, to wait for messages to be polled.
      Returns:
      A mapping of PubSub topic partitions to lists of PubSub messages retrieved from Kafka.

    • hasAnySubscription

      public boolean hasAnySubscription()
      Description copied from interface: PubSubConsumerAdapter
      Checks if the consumer has any active topic-partition subscriptions.
      Specified by:
      hasAnySubscription in interface PubSubConsumerAdapter
      Returns:
      true if the consumer is subscribed to one or more topic partitions, false otherwise.

    • hasSubscription

      public boolean hasSubscription(PubSubTopicPartition pubSubTopicPartition)
      Description copied from interface: PubSubConsumerAdapter
      Checks if the consumer is currently subscribed to the specified PubSub topic-partition.
      Specified by:
      hasSubscription in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic-partition to check for subscription.
      Returns:
      true if the consumer is subscribed to the given topic-partition, false otherwise.

    • getOffsets

      public Map<PubSubTopicPartition,Long> getOffsets()

    • pause

      public void pause(PubSubTopicPartition pubSubTopicPartition)
      Description copied from interface: PubSubConsumerAdapter
      Pauses message consumption for the specified PubSub topic-partition. If the partition was not previously subscribed, this method is a no-op.
      Specified by:
      pause in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic-partition for which to pause message consumption.

    • resume

      public void resume(PubSubTopicPartition pubSubTopicPartition)
      Description copied from interface: PubSubConsumerAdapter
      Resumes message consumption for the specified PubSub topic-partition. If the partition was not previously paused or if they were not subscribed at all, this method is a no-op.
      Specified by:
      resume in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic partition for which to resume message consumption.

    • getAssignment

      public Set<PubSubTopicPartition> getAssignment()
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the set of PubSub topic-partitions currently assigned to the consumer.
      Specified by:
      getAssignment in interface PubSubConsumerAdapter
      Returns:
      A set of PubSub topic-partitions representing the current assignment of the consumer.

    • offsetForTime

      public Long offsetForTime(PubSubTopicPartition pubSubTopicPartition, long timestamp, Duration timeout)
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the offset of the first message with a timestamp greater than or equal to the target timestamp for the specified PubSub topic-partition. If no such message is found, null will be returned for the partition.
      Specified by:
      offsetForTime in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic-partition for which to fetch the offset.
      timestamp - The target timestamp to search for in milliseconds since the Unix epoch.
      timeout - The maximum duration to wait for the operation to complete.
      Returns:
      The offset of the first message with a timestamp greater than or equal to the target timestamp, or null if no such message is found for the partition.

    • offsetForTime

      public Long offsetForTime(PubSubTopicPartition pubSubTopicPartition, long timestamp)
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the offset of the first message with a timestamp greater than or equal to the target timestamp for the specified PubSub topic-partition. If no such message is found, null will be returned for the partition.
      Specified by:
      offsetForTime in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic-partition for which to fetch the offset.
      timestamp - The target timestamp to search for in milliseconds since the Unix epoch.
      Returns:
      The offset of the first message with a timestamp greater than or equal to the target timestamp, or null if no such message is found for the partition.

    • beginningOffset

      public Long beginningOffset(PubSubTopicPartition partition, Duration timeout)
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the beginning offset for the specified PubSub topic-partition.
      Specified by:
      beginningOffset in interface PubSubConsumerAdapter
      Parameters:
      partition - The PubSub topic-partition for which to fetch the beginning offset.
      timeout - The maximum duration to wait for the operation to complete.
      Returns:
      The beginning offset of the specified topic-partition. If topic-partition exists but has no messages, the offset will be 0.

    • endOffsets

      public Map<PubSubTopicPartition,Long> endOffsets(Collection<PubSubTopicPartition> partitions, Duration timeout)
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the end offsets for a collection of PubSub topic-partitions. The end offset represents the highest offset available in each specified partition, i.e., offset of the last message + 1. If there are no messages in a partition, the end offset will be 0.
      Specified by:
      endOffsets in interface PubSubConsumerAdapter
      Parameters:
      partitions - A collection of PubSub topic-partitions for which to fetch the end offsets.
      timeout - The maximum duration to wait for the operation to complete.
      Returns:
      A mapping of PubSub topic partitions to their respective end offsets, or an empty map if the offsets cannot be determined.

    • endOffset

      public Long endOffset(PubSubTopicPartition pubSubTopicPartition)
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the end offset for the specified PubSub topic-partition. The end offset represents the highest offset available in each specified partition, i.e., offset of the last message + 1. If there are no messages in a partition, the end offset will be 0.
      Specified by:
      endOffset in interface PubSubConsumerAdapter
      Parameters:
      pubSubTopicPartition - The PubSub topic partition for which to fetch the end offset.
      Returns:
      The end offset of the specified topic partition.

    • partitionsFor

      public List<PubSubTopicPartitionInfo> partitionsFor(PubSubTopic topic)
      Description copied from interface: PubSubConsumerAdapter
      Retrieves the list of partitions associated with a given Pub-Sub topic.
      Specified by:
      partitionsFor in interface PubSubConsumerAdapter
      Parameters:
      topic - The Pub-Sub topic for which partition information is requested.
      Returns:
      A list of PubSubTopicPartitionInfo representing the partitions of the topic, or null if the topic does not exist.

    • setMockInMemoryAdminAdapter

      public void setMockInMemoryAdminAdapter(MockInMemoryAdminAdapter adminAdapter)