Search in sources :

Example 26 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StreamThreadTest method shouldChangeStateInRebalanceListener.

@Test
public void shouldChangeStateInRebalanceListener() {
    final StreamThread thread = createStreamThread(CLIENT_ID, config, false);
    final StateListenerStub stateListener = new StateListenerStub();
    thread.setStateListener(stateListener);
    assertEquals(thread.state(), StreamThread.State.CREATED);
    final ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener();
    final List<TopicPartition> revokedPartitions;
    final List<TopicPartition> assignedPartitions;
    // revoke nothing
    thread.setState(StreamThread.State.STARTING);
    revokedPartitions = Collections.emptyList();
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    assertEquals(thread.state(), StreamThread.State.PARTITIONS_REVOKED);
    // assign single partition
    assignedPartitions = Collections.singletonList(t1p1);
    final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.mainConsumer();
    mockConsumer.assign(assignedPartitions);
    mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    thread.runOnce();
    assertEquals(thread.state(), StreamThread.State.RUNNING);
    Assert.assertEquals(4, stateListener.numChanges);
    Assert.assertEquals(StreamThread.State.PARTITIONS_ASSIGNED, stateListener.oldState);
    thread.shutdown();
    assertSame(StreamThread.State.PENDING_SHUTDOWN, thread.state());
}
Also used : TopicPartition(org.apache.kafka.common.TopicPartition) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) ConsumerRebalanceListener(org.apache.kafka.clients.consumer.ConsumerRebalanceListener) Test(org.junit.Test)

Example 27 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class StreamTaskTest method createDisconnectedTask.

private StreamTask createDisconnectedTask(final StreamsConfig config) {
    final MockKeyValueStore stateStore = new MockKeyValueStore(storeName, false);
    final ProcessorTopology topology = ProcessorTopologyFactories.with(asList(source1, source2), mkMap(mkEntry(topic1, source1), mkEntry(topic2, source2)), singletonList(stateStore), emptyMap());
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {

        @Override
        public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition> partitions) {
            throw new TimeoutException("KABOOM!");
        }
    };
    final InternalProcessorContext context = new ProcessorContextImpl(taskId, config, stateManager, streamsMetrics, null);
    return new StreamTask(taskId, partitions, topology, consumer, new TopologyConfig(null, config, new Properties()).getTaskConfig(), streamsMetrics, stateDirectory, cache, time, stateManager, recordCollector, context, logContext);
}
Also used : Utils.mkSet(org.apache.kafka.common.utils.Utils.mkSet) Set(java.util.Set) HashSet(java.util.HashSet) MockKeyValueStore(org.apache.kafka.test.MockKeyValueStore) Utils.mkProperties(org.apache.kafka.common.utils.Utils.mkProperties) Properties(java.util.Properties) TopicPartition(org.apache.kafka.common.TopicPartition) OffsetAndMetadata(org.apache.kafka.clients.consumer.OffsetAndMetadata) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) TopologyConfig(org.apache.kafka.streams.processor.internals.namedtopology.TopologyConfig) TimeoutException(org.apache.kafka.common.errors.TimeoutException)

Example 28 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project kafka by apache.

the class GlobalStreamThreadTest method before.

@Before
public void before() {
    final MaterializedInternal<Object, Object, KeyValueStore<Bytes, byte[]>> materialized = new MaterializedInternal<>(Materialized.with(null, null), new InternalNameProvider() {

        @Override
        public String newProcessorName(final String prefix) {
            return "processorName";
        }

        @Override
        public String newStoreName(final String prefix) {
            return GLOBAL_STORE_NAME;
        }
    }, "store-");
    final ProcessorSupplier<Object, Object, Void, Void> processorSupplier = () -> new ContextualProcessor<Object, Object, Void, Void>() {

        @Override
        public void process(final Record<Object, Object> record) {
        }
    };
    builder.addGlobalStore(new TimestampedKeyValueStoreMaterializer<>(materialized).materialize().withLoggingDisabled(), "sourceName", null, null, null, GLOBAL_STORE_TOPIC_NAME, "processorName", processorSupplier);
    baseDirectoryName = TestUtils.tempDirectory().getAbsolutePath();
    final HashMap<String, Object> properties = new HashMap<>();
    properties.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "blah");
    properties.put(StreamsConfig.APPLICATION_ID_CONFIG, "testAppId");
    properties.put(StreamsConfig.STATE_DIR_CONFIG, baseDirectoryName);
    properties.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArraySerde.class.getName());
    properties.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArraySerde.class.getName());
    config = new StreamsConfig(properties);
    globalStreamThread = new GlobalStreamThread(builder.rewriteTopology(config).buildGlobalStateTopology(), config, mockConsumer, new StateDirectory(config, time, true, false), 0, new StreamsMetricsImpl(new Metrics(), "test-client", StreamsConfig.METRICS_LATEST, time), time, "clientId", stateRestoreListener, e -> {
    });
}
Also used : CoreMatchers.is(org.hamcrest.CoreMatchers.is) StreamsConfig(org.apache.kafka.streams.StreamsConfig) MockTime(org.apache.kafka.common.utils.MockTime) InternalNameProvider(org.apache.kafka.streams.kstream.internals.InternalNameProvider) IsInstanceOf.instanceOf(org.hamcrest.core.IsInstanceOf.instanceOf) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) CoreMatchers.equalTo(org.hamcrest.CoreMatchers.equalTo) HashMap(java.util.HashMap) StreamsException(org.apache.kafka.streams.errors.StreamsException) OffsetResetStrategy(org.apache.kafka.clients.consumer.OffsetResetStrategy) DEAD(org.apache.kafka.streams.processor.internals.GlobalStreamThread.State.DEAD) ContextualProcessor(org.apache.kafka.streams.processor.api.ContextualProcessor) ProcessorSupplier(org.apache.kafka.streams.processor.api.ProcessorSupplier) KeyValueStore(org.apache.kafka.streams.state.KeyValueStore) Serdes(org.apache.kafka.common.serialization.Serdes) Record(org.apache.kafka.streams.processor.api.Record) StreamsMetricsImpl(org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl) Assert.fail(org.junit.Assert.fail) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) Before(org.junit.Before) TopicPartition(org.apache.kafka.common.TopicPartition) RUNNING(org.apache.kafka.streams.processor.internals.GlobalStreamThread.State.RUNNING) MockStateRestoreListener(org.apache.kafka.test.MockStateRestoreListener) TestUtils(org.apache.kafka.test.TestUtils) Set(java.util.Set) Assert.assertTrue(org.junit.Assert.assertTrue) Test(org.junit.Test) PartitionInfo(org.apache.kafka.common.PartitionInfo) MaterializedInternal(org.apache.kafka.streams.kstream.internals.MaterializedInternal) File(java.io.File) Bytes(org.apache.kafka.common.utils.Bytes) ConsumerRecordUtil.record(org.apache.kafka.streams.processor.internals.testutil.ConsumerRecordUtil.record) List(java.util.List) Metrics(org.apache.kafka.common.metrics.Metrics) TimestampedKeyValueStoreMaterializer(org.apache.kafka.streams.kstream.internals.TimestampedKeyValueStoreMaterializer) StateStore(org.apache.kafka.streams.processor.StateStore) Assert.assertFalse(org.junit.Assert.assertFalse) Materialized(org.apache.kafka.streams.kstream.Materialized) InvalidOffsetException(org.apache.kafka.clients.consumer.InvalidOffsetException) Node(org.apache.kafka.common.Node) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) HashMap(java.util.HashMap) KeyValueStore(org.apache.kafka.streams.state.KeyValueStore) InternalNameProvider(org.apache.kafka.streams.kstream.internals.InternalNameProvider) Metrics(org.apache.kafka.common.metrics.Metrics) Record(org.apache.kafka.streams.processor.api.Record) ContextualProcessor(org.apache.kafka.streams.processor.api.ContextualProcessor) MaterializedInternal(org.apache.kafka.streams.kstream.internals.MaterializedInternal) StreamsMetricsImpl(org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl) StreamsConfig(org.apache.kafka.streams.StreamsConfig) Before(org.junit.Before)

Example 29 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project beam by apache.

the class KafkaIOTest method mkMockConsumer.

// Update mock consumer with records distributed among the given topics, each with given number
// of partitions. Records are assigned in round-robin order among the partitions.
private static MockConsumer<byte[], byte[]> mkMockConsumer(List<String> topics, int partitionsPerTopic, int numElements, OffsetResetStrategy offsetResetStrategy, Map<String, Object> config, SerializableFunction<Integer, byte[]> keyFunction, SerializableFunction<Integer, byte[]> valueFunction) {
    final List<TopicPartition> partitions = new ArrayList<>();
    final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> records = new HashMap<>();
    Map<String, List<PartitionInfo>> partitionMap = new HashMap<>();
    for (String topic : topics) {
        List<PartitionInfo> partIds = new ArrayList<>(partitionsPerTopic);
        for (int i = 0; i < partitionsPerTopic; i++) {
            TopicPartition tp = new TopicPartition(topic, i);
            partitions.add(tp);
            partIds.add(new PartitionInfo(topic, i, null, null, null));
            records.put(tp, new ArrayList<>());
        }
        partitionMap.put(topic, partIds);
    }
    int numPartitions = partitions.size();
    final long[] offsets = new long[numPartitions];
    long timestampStartMillis = (Long) config.getOrDefault(TIMESTAMP_START_MILLIS_CONFIG, LOG_APPEND_START_TIME.getMillis());
    TimestampType timestampType = TimestampType.forName((String) config.getOrDefault(TIMESTAMP_TYPE_CONFIG, TimestampType.LOG_APPEND_TIME.toString()));
    for (int i = 0; i < numElements; i++) {
        int pIdx = i % numPartitions;
        TopicPartition tp = partitions.get(pIdx);
        byte[] key = keyFunction.apply(i);
        byte[] value = valueFunction.apply(i);
        records.get(tp).add(new ConsumerRecord<>(tp.topic(), tp.partition(), offsets[pIdx]++, timestampStartMillis + Duration.standardSeconds(i).getMillis(), timestampType, 0, key.length, value.length, key, value));
    }
    // This is updated when reader assigns partitions.
    final AtomicReference<List<TopicPartition>> assignedPartitions = new AtomicReference<>(Collections.<TopicPartition>emptyList());
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(offsetResetStrategy) {

        @Override
        public synchronized void assign(final Collection<TopicPartition> assigned) {
            super.assign(assigned);
            assignedPartitions.set(ImmutableList.copyOf(assigned));
            for (TopicPartition tp : assigned) {
                updateBeginningOffsets(ImmutableMap.of(tp, 0L));
                updateEndOffsets(ImmutableMap.of(tp, (long) records.get(tp).size()));
            }
        }

        // Override offsetsForTimes() in order to look up the offsets by timestamp.
        @Override
        public synchronized Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(Map<TopicPartition, Long> timestampsToSearch) {
            return timestampsToSearch.entrySet().stream().map(e -> {
                // In test scope, timestamp == offset.
                long maxOffset = offsets[partitions.indexOf(e.getKey())];
                long offset = e.getValue();
                OffsetAndTimestamp value = (offset >= maxOffset) ? null : new OffsetAndTimestamp(offset, offset);
                return new SimpleEntry<>(e.getKey(), value);
            }).collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
        }
    };
    for (String topic : topics) {
        consumer.updatePartitions(topic, partitionMap.get(topic));
    }
    // MockConsumer does not maintain any relationship between partition seek position and the
    // records added. e.g. if we add 10 records to a partition and then seek to end of the
    // partition, MockConsumer is still going to return the 10 records in next poll. It is
    // our responsibility to make sure currently enqueued records sync with partition offsets.
    // The following task will be called inside each invocation to MockConsumer.poll().
    // We enqueue only the records with the offset >= partition's current position.
    Runnable recordEnqueueTask = new Runnable() {

        @Override
        public void run() {
            // add all the records with offset >= current partition position.
            int recordsAdded = 0;
            for (TopicPartition tp : assignedPartitions.get()) {
                long curPos = consumer.position(tp);
                for (ConsumerRecord<byte[], byte[]> r : records.get(tp)) {
                    if (r.offset() >= curPos) {
                        consumer.addRecord(r);
                        recordsAdded++;
                    }
                }
            }
            if (recordsAdded == 0) {
                if (config.get("inject.error.at.eof") != null) {
                    consumer.setException(new KafkaException("Injected error in consumer.poll()"));
                }
                // MockConsumer.poll(timeout) does not actually wait even when there aren't any
                // records.
                // Add a small wait here in order to avoid busy looping in the reader.
                Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
            // TODO: BEAM-4086: testUnboundedSourceWithoutBoundedWrapper() occasionally hangs
            // without this wait. Need to look into it.
            }
            consumer.schedulePollTask(this);
        }
    };
    consumer.schedulePollTask(recordEnqueueTask);
    return consumer;
}
Also used : Count(org.apache.beam.sdk.transforms.Count) MetricName(org.apache.beam.sdk.metrics.MetricName) Arrays(java.util.Arrays) PipelineExecutionException(org.apache.beam.sdk.Pipeline.PipelineExecutionException) SchemaRegistryClient(io.confluent.kafka.schemaregistry.client.SchemaRegistryClient) CoderUtils(org.apache.beam.sdk.util.CoderUtils) KafkaAvroSerializer(io.confluent.kafka.serializers.KafkaAvroSerializer) UnboundedSource(org.apache.beam.sdk.io.UnboundedSource) KafkaException(org.apache.kafka.common.KafkaException) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) OffsetResetStrategy(org.apache.kafka.clients.consumer.OffsetResetStrategy) Map(java.util.Map) Window(org.apache.beam.sdk.transforms.windowing.Window) Uninterruptibles(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles) TimestampType(org.apache.kafka.common.record.TimestampType) MockSchemaRegistry(io.confluent.kafka.schemaregistry.testutil.MockSchemaRegistry) MetricResult(org.apache.beam.sdk.metrics.MetricResult) UnboundedReader(org.apache.beam.sdk.io.UnboundedSource.UnboundedReader) ConsumerConfig(org.apache.kafka.clients.consumer.ConsumerConfig) PartitionInfo(org.apache.kafka.common.PartitionInfo) OffsetAndTimestamp(org.apache.kafka.clients.consumer.OffsetAndTimestamp) StandardCharsets(java.nio.charset.StandardCharsets) Executors(java.util.concurrent.Executors) Matchers.instanceOf(org.hamcrest.Matchers.instanceOf) ConsumerRecord(org.apache.kafka.clients.consumer.ConsumerRecord) MetricQueryResults(org.apache.beam.sdk.metrics.MetricQueryResults) Matchers.greaterThan(org.hamcrest.Matchers.greaterThan) DisplayDataMatchers.hasDisplayItem(org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem) Matchers.containsString(org.hamcrest.Matchers.containsString) KV(org.apache.beam.sdk.values.KV) Duration(org.joda.time.Duration) RunWith(org.junit.runner.RunWith) RecordHeader(org.apache.kafka.common.header.internals.RecordHeader) ArrayList(java.util.ArrayList) SinkMetrics(org.apache.beam.sdk.metrics.SinkMetrics) Read(org.apache.beam.sdk.io.Read) Distinct(org.apache.beam.sdk.transforms.Distinct) ThreadLocalRandom(java.util.concurrent.ThreadLocalRandom) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) IntegerSerializer(org.apache.kafka.common.serialization.IntegerSerializer) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) Deserializer(org.apache.kafka.common.serialization.Deserializer) DoFn(org.apache.beam.sdk.transforms.DoFn) ByteArrayDeserializer(org.apache.kafka.common.serialization.ByteArrayDeserializer) ThrowableMessageMatcher.hasMessage(org.junit.internal.matchers.ThrowableMessageMatcher.hasMessage) PAssert(org.apache.beam.sdk.testing.PAssert) Producer(org.apache.kafka.clients.producer.Producer) Assert.assertTrue(org.junit.Assert.assertTrue) IOException(java.io.IOException) Test(org.junit.Test) AbstractKafkaAvroSerDeConfig(io.confluent.kafka.serializers.AbstractKafkaAvroSerDeConfig) MetricResultsMatchers.attemptedMetricsResult(org.apache.beam.sdk.metrics.MetricResultsMatchers.attemptedMetricsResult) Matchers.hasItem(org.hamcrest.Matchers.hasItem) Assert.assertNull(org.junit.Assert.assertNull) IsIterableWithSize(org.hamcrest.collection.IsIterableWithSize) Serializer(org.apache.kafka.common.serialization.Serializer) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) VarIntCoder(org.apache.beam.sdk.coders.VarIntCoder) Assert.assertEquals(org.junit.Assert.assertEquals) MockProducer(org.apache.kafka.clients.producer.MockProducer) SourceMetrics(org.apache.beam.sdk.metrics.SourceMetrics) Matchers.isA(org.hamcrest.Matchers.isA) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) PipelineResult(org.apache.beam.sdk.PipelineResult) LoggerFactory(org.slf4j.LoggerFactory) SerializableFunction(org.apache.beam.sdk.transforms.SerializableFunction) ByteBuffer(java.nio.ByteBuffer) PCollectionList(org.apache.beam.sdk.values.PCollectionList) ThrowableCauseMatcher.hasCause(org.junit.internal.matchers.ThrowableCauseMatcher.hasCause) Method(java.lang.reflect.Method) Flatten(org.apache.beam.sdk.transforms.Flatten) MapElements(org.apache.beam.sdk.transforms.MapElements) Min(org.apache.beam.sdk.transforms.Min) Consumer(org.apache.kafka.clients.consumer.Consumer) TopicPartition(org.apache.kafka.common.TopicPartition) Collection(java.util.Collection) BigEndianIntegerCoder(org.apache.beam.sdk.coders.BigEndianIntegerCoder) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) LongDeserializer(org.apache.kafka.common.serialization.LongDeserializer) VarLongCoder(org.apache.beam.sdk.coders.VarLongCoder) LongSerializer(org.apache.kafka.common.serialization.LongSerializer) Collectors(java.util.stream.Collectors) List(java.util.List) Max(org.apache.beam.sdk.transforms.Max) ParDo(org.apache.beam.sdk.transforms.ParDo) Header(org.apache.kafka.common.header.Header) TypeDescriptors(org.apache.beam.sdk.values.TypeDescriptors) Optional(java.util.Optional) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Assume.assumeTrue(org.junit.Assume.assumeTrue) Values(org.apache.beam.sdk.transforms.Values) MetricNameFilter(org.apache.beam.sdk.metrics.MetricNameFilter) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) Headers(org.apache.kafka.common.header.Headers) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) HashMap(java.util.HashMap) AvroGeneratedUser(org.apache.beam.sdk.io.AvroGeneratedUser) PipelineOptionsFactory(org.apache.beam.sdk.options.PipelineOptionsFactory) AtomicReference(java.util.concurrent.atomic.AtomicReference) ConcurrentMap(java.util.concurrent.ConcurrentMap) SimpleEntry(java.util.AbstractMap.SimpleEntry) ProducerConfig(org.apache.kafka.clients.producer.ProducerConfig) ExpectedException(org.junit.rules.ExpectedException) ExecutorService(java.util.concurrent.ExecutorService) Nullable(org.checkerframework.checker.nullness.qual.Nullable) Utils(org.apache.kafka.common.utils.Utils) DisplayData(org.apache.beam.sdk.transforms.display.DisplayData) GenericRecord(org.apache.avro.generic.GenericRecord) Logger(org.slf4j.Logger) BigEndianLongCoder(org.apache.beam.sdk.coders.BigEndianLongCoder) Assert.assertNotNull(org.junit.Assert.assertNotNull) FixedWindows(org.apache.beam.sdk.transforms.windowing.FixedWindows) Lists(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists) JUnit4(org.junit.runners.JUnit4) PCollection(org.apache.beam.sdk.values.PCollection) MetricsFilter(org.apache.beam.sdk.metrics.MetricsFilter) TimeUnit(java.util.concurrent.TimeUnit) Rule(org.junit.Rule) Ignore(org.junit.Ignore) ConfluentSchemaRegistryDeserializerProviderTest.mockDeserializerProvider(org.apache.beam.sdk.io.kafka.ConfluentSchemaRegistryDeserializerProviderTest.mockDeserializerProvider) Instant(org.joda.time.Instant) IntegerDeserializer(org.apache.kafka.common.serialization.IntegerDeserializer) Comparator(java.util.Comparator) Collections(java.util.Collections) IsIterableContainingInAnyOrder(org.hamcrest.collection.IsIterableContainingInAnyOrder) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) Matchers.containsString(org.hamcrest.Matchers.containsString) TimestampType(org.apache.kafka.common.record.TimestampType) ArrayList(java.util.ArrayList) PCollectionList(org.apache.beam.sdk.values.PCollectionList) List(java.util.List) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) PartitionInfo(org.apache.kafka.common.PartitionInfo) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) OffsetAndTimestamp(org.apache.kafka.clients.consumer.OffsetAndTimestamp) AtomicReference(java.util.concurrent.atomic.AtomicReference) TopicPartition(org.apache.kafka.common.TopicPartition) Collection(java.util.Collection) PCollection(org.apache.beam.sdk.values.PCollection) KafkaException(org.apache.kafka.common.KafkaException) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap) ConcurrentMap(java.util.concurrent.ConcurrentMap)

Example 30 with MockConsumer

use of org.apache.kafka.clients.consumer.MockConsumer in project beam by apache.

the class KafkaTestTable method mkMockConsumer.

private MockConsumer<byte[], byte[]> mkMockConsumer(Map<String, Object> config) {
    OffsetResetStrategy offsetResetStrategy = OffsetResetStrategy.EARLIEST;
    final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> kafkaRecords = new HashMap<>();
    Map<String, List<PartitionInfo>> partitionInfoMap = new HashMap<>();
    Map<String, List<TopicPartition>> partitionMap = new HashMap<>();
    // Create Topic Paritions
    for (String topic : this.getTopics()) {
        List<PartitionInfo> partIds = new ArrayList<>(partitionsPerTopic);
        List<TopicPartition> topicParitions = new ArrayList<>(partitionsPerTopic);
        for (int i = 0; i < partitionsPerTopic; i++) {
            TopicPartition tp = new TopicPartition(topic, i);
            topicParitions.add(tp);
            partIds.add(new PartitionInfo(topic, i, null, null, null));
            kafkaRecords.put(tp, new ArrayList<>());
        }
        partitionInfoMap.put(topic, partIds);
        partitionMap.put(topic, topicParitions);
    }
    TimestampType timestampType = TimestampType.forName((String) config.getOrDefault(TIMESTAMP_TYPE_CONFIG, TimestampType.LOG_APPEND_TIME.toString()));
    for (KafkaTestRecord record : this.records) {
        int partitionIndex = record.getKey().hashCode() % partitionsPerTopic;
        TopicPartition tp = partitionMap.get(record.getTopic()).get(partitionIndex);
        byte[] key = record.getKey().getBytes(UTF_8);
        byte[] value = record.getValue().toByteArray();
        kafkaRecords.get(tp).add(new ConsumerRecord<>(tp.topic(), tp.partition(), kafkaRecords.get(tp).size(), record.getTimeStamp(), timestampType, 0, key.length, value.length, key, value));
    }
    // This is updated when reader assigns partitions.
    final AtomicReference<List<TopicPartition>> assignedPartitions = new AtomicReference<>(Collections.emptyList());
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(offsetResetStrategy) {

        @Override
        public synchronized void assign(final Collection<TopicPartition> assigned) {
            Collection<TopicPartition> realPartitions = assigned.stream().map(part -> partitionMap.get(part.topic()).get(part.partition())).collect(Collectors.toList());
            super.assign(realPartitions);
            assignedPartitions.set(ImmutableList.copyOf(realPartitions));
            for (TopicPartition tp : realPartitions) {
                updateBeginningOffsets(ImmutableMap.of(tp, 0L));
                updateEndOffsets(ImmutableMap.of(tp, (long) kafkaRecords.get(tp).size()));
            }
        }

        // Override offsetsForTimes() in order to look up the offsets by timestamp.
        @Override
        public synchronized Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(Map<TopicPartition, Long> timestampsToSearch) {
            return timestampsToSearch.entrySet().stream().map(e -> {
                // In test scope, timestamp == offset. ????
                long maxOffset = kafkaRecords.get(e.getKey()).size();
                long offset = e.getValue();
                OffsetAndTimestamp value = (offset >= maxOffset) ? null : new OffsetAndTimestamp(offset, offset);
                return new AbstractMap.SimpleEntry<>(e.getKey(), value);
            }).collect(Collectors.toMap(AbstractMap.SimpleEntry::getKey, AbstractMap.SimpleEntry::getValue));
        }
    };
    for (String topic : getTopics()) {
        consumer.updatePartitions(topic, partitionInfoMap.get(topic));
    }
    Runnable recordEnqueueTask = new Runnable() {

        @Override
        public void run() {
            // add all the records with offset >= current partition position.
            int recordsAdded = 0;
            for (TopicPartition tp : assignedPartitions.get()) {
                long curPos = consumer.position(tp);
                for (ConsumerRecord<byte[], byte[]> r : kafkaRecords.get(tp)) {
                    if (r.offset() >= curPos) {
                        consumer.addRecord(r);
                        recordsAdded++;
                    }
                }
            }
            if (recordsAdded == 0) {
                if (config.get("inject.error.at.eof") != null) {
                    consumer.setException(new KafkaException("Injected error in consumer.poll()"));
                }
                // MockConsumer.poll(timeout) does not actually wait even when there aren't any
                // records.
                // Add a small wait here in order to avoid busy looping in the reader.
                Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
            }
            consumer.schedulePollTask(this);
        }
    };
    consumer.schedulePollTask(recordEnqueueTask);
    return consumer;
}
Also used : ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) KafkaException(org.apache.kafka.common.KafkaException) HashMap(java.util.HashMap) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) OffsetResetStrategy(org.apache.kafka.clients.consumer.OffsetResetStrategy) AtomicReference(java.util.concurrent.atomic.AtomicReference) ArrayList(java.util.ArrayList) KafkaIO(org.apache.beam.sdk.io.kafka.KafkaIO) PTransform(org.apache.beam.sdk.transforms.PTransform) Map(java.util.Map) Uninterruptibles(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles) Row(org.apache.beam.sdk.values.Row) TimestampType(org.apache.kafka.common.record.TimestampType) TopicPartition(org.apache.kafka.common.TopicPartition) KafkaRecord(org.apache.beam.sdk.io.kafka.KafkaRecord) UTF_8(java.nio.charset.StandardCharsets.UTF_8) Collection(java.util.Collection) PartitionInfo(org.apache.kafka.common.PartitionInfo) PCollection(org.apache.beam.sdk.values.PCollection) OffsetAndTimestamp(org.apache.kafka.clients.consumer.OffsetAndTimestamp) Collectors(java.util.stream.Collectors) Schema(org.apache.beam.sdk.schemas.Schema) TimeUnit(java.util.concurrent.TimeUnit) AbstractMap(java.util.AbstractMap) List(java.util.List) ConsumerRecord(org.apache.kafka.clients.consumer.ConsumerRecord) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Collections(java.util.Collections) OffsetResetStrategy(org.apache.kafka.clients.consumer.OffsetResetStrategy) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) AbstractMap(java.util.AbstractMap) TimestampType(org.apache.kafka.common.record.TimestampType) ArrayList(java.util.ArrayList) List(java.util.List) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) PartitionInfo(org.apache.kafka.common.PartitionInfo) MockConsumer(org.apache.kafka.clients.consumer.MockConsumer) OffsetAndTimestamp(org.apache.kafka.clients.consumer.OffsetAndTimestamp) AtomicReference(java.util.concurrent.atomic.AtomicReference) TopicPartition(org.apache.kafka.common.TopicPartition) Collection(java.util.Collection) PCollection(org.apache.beam.sdk.values.PCollection) KafkaException(org.apache.kafka.common.KafkaException) HashMap(java.util.HashMap) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map) AbstractMap(java.util.AbstractMap)

Aggregations

MockConsumer (org.apache.kafka.clients.consumer.MockConsumer)56 Test (org.junit.Test)46 TopicPartition (org.apache.kafka.common.TopicPartition)44 HashMap (java.util.HashMap)27 TaskId (org.apache.kafka.streams.processor.TaskId)27 Set (java.util.Set)24 ArrayList (java.util.ArrayList)20 StreamsConfig (org.apache.kafka.streams.StreamsConfig)20 PartitionInfo (org.apache.kafka.common.PartitionInfo)18 HashSet (java.util.HashSet)17 Utils.mkSet (org.apache.kafka.common.utils.Utils.mkSet)15 Map (java.util.Map)10 Properties (java.util.Properties)10 StreamsException (org.apache.kafka.streams.errors.StreamsException)10 Collections.emptySet (java.util.Collections.emptySet)9 InternalStreamsBuilderTest (org.apache.kafka.streams.kstream.internals.InternalStreamsBuilderTest)9 KafkaException (org.apache.kafka.common.KafkaException)8 TimeoutException (org.apache.kafka.common.errors.TimeoutException)8 List (java.util.List)7 AtomicBoolean (java.util.concurrent.atomic.AtomicBoolean)7