Examples with Subscription org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription used on opensource projects

Example 21 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class RoundRobinAssignorTest method testTwoConsumersOneTopicTwoPartitions.

@Test
public void testTwoConsumersOneTopicTwoPartitions() {
    String topic = "topic";
    String consumer1 = "consumer1";
    String consumer2 = "consumer2";
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 2);
    Map<String, Subscription> consumers = new HashMap<>();
    consumers.put(consumer1, new Subscription(topics(topic)));
    consumers.put(consumer2, new Subscription(topics(topic)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
    assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
    assertEquals(partitions(tp(topic, 1)), assignment.get(consumer2));
}

Example 22 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class RoundRobinAssignorTest method testTwoConsumersOneTopicOnePartition.

@Test
public void testTwoConsumersOneTopicOnePartition() {
    String topic = "topic";
    String consumer1 = "consumer1";
    String consumer2 = "consumer2";
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 1);
    Map<String, Subscription> consumers = new HashMap<>();
    consumers.put(consumer1, new Subscription(topics(topic)));
    consumers.put(consumer2, new Subscription(topics(topic)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
    assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
    assertEquals(Collections.<TopicPartition>emptyList(), assignment.get(consumer2));
}

Example 23 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testNewSubscription.

@Test
public void testNewSubscription() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 1; i < 5; i++) partitionsPerTopic.put(getTopicName(i, 5), 1);
    Map<String, Subscription> subscriptions = new HashMap<>();
    for (int i = 0; i < 3; i++) {
        List<String> topics = new ArrayList<String>();
        for (int j = i; j <= 3 * i - 2; j++) topics.add(getTopicName(j, 5));
        subscriptions.put(getConsumerName(i, 3), new Subscription(topics));
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    subscriptions.get(getConsumerName(0, 3)).topics().add(getTopicName(1, 5));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(assignor.isSticky());
}

Example 24 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testOneConsumerNoTopic.

@Test
public void testOneConsumerNoTopic() {
    String consumerId = "consumer";
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    Map<String, Subscription> subscriptions = Collections.singletonMap(consumerId, new Subscription(Collections.<String>emptyList()));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(Collections.singleton(consumerId), assignment.keySet());
    assertTrue(assignment.get(consumerId).isEmpty());
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(isFullyBalanced(assignment));
}

Example 25 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testReassignmentAfterOneConsumerLeaves.

@Test
public void testReassignmentAfterOneConsumerLeaves() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 1; i < 20; i++) partitionsPerTopic.put(getTopicName(i, 20), i);
    Map<String, Subscription> subscriptions = new HashMap<>();
    for (int i = 1; i < 20; i++) {
        List<String> topics = new ArrayList<String>();
        for (int j = 1; j <= i; j++) topics.add(getTopicName(j, 20));
        subscriptions.put(getConsumerName(i, 20), new Subscription(topics));
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    for (int i = 1; i < 20; i++) {
        String consumer = getConsumerName(i, 20);
        subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
    }
    subscriptions.remove("consumer10");
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(assignor.isSticky());
}