Java MockProducer类代码示例

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
private static void verifyProducerRecords(MockProducer<Integer, Long> mockProducer,
                                          String topic, int numElements, boolean keyIsAbsent) {

  // verify that appropriate messages are written to kafka
  List<ProducerRecord<Integer, Long>> sent = mockProducer.history();

  // sort by values
  Collections.sort(sent, new Comparator<ProducerRecord<Integer, Long>>() {
    @Override
    public int compare(ProducerRecord<Integer, Long> o1, ProducerRecord<Integer, Long> o2) {
      return Long.compare(o1.value(), o2.value());
    }
  });

  for (int i = 0; i < numElements; i++) {
    ProducerRecord<Integer, Long> record = sent.get(i);
    assertEquals(topic, record.topic());
    if (keyIsAbsent) {
      assertNull(record.key());
    } else {
      assertEquals(i, record.key().intValue());
    }
    assertEquals(i, record.value().longValue());
  }
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@Test
public void testStreamPartitioner() {

    final RecordCollectorImpl collector = new RecordCollectorImpl(
            new MockProducer<>(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer),
            "RecordCollectorTest-TestStreamPartitioner");

    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "9", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "27", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "81", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "243", "0", null, stringSerializer, stringSerializer, streamPartitioner);

    collector.send("topic1", "28", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "82", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "244", "0", null, stringSerializer, stringSerializer, streamPartitioner);

    collector.send("topic1", "245", "0", null, stringSerializer, stringSerializer, streamPartitioner);

    final Map<TopicPartition, Long> offsets = collector.offsets();

    assertEquals((Long) 4L, offsets.get(new TopicPartition("topic1", 0)));
    assertEquals((Long) 2L, offsets.get(new TopicPartition("topic1", 1)));
    assertEquals((Long) 0L, offsets.get(new TopicPartition("topic1", 2)));
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
MockProducerWrapper() {
  producerKey = String.valueOf(ThreadLocalRandom.current().nextLong());
  mockProducer = new MockProducer<Integer, Long>(
    false, // disable synchronous completion of send. see ProducerSendCompletionThread below.
    new IntegerSerializer(),
    new LongSerializer()) {

    // override flush() so that it does not complete all the waiting sends, giving a chance to
    // ProducerCompletionThread to inject errors.

    @Override
    public void flush() {
      while (completeNext()) {
        // there are some uncompleted records. let the completion thread handle them.
        try {
          Thread.sleep(10);
        } catch (InterruptedException e) {
          // ok to retry.
        }
      }
    }
  };

  // Add the producer to the global map so that producer factory function can access it.
  assertNull(MOCK_PRODUCER_MAP.putIfAbsent(producerKey, mockProducer));
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
private MockProducer createProducer() {
    Collection<PartitionInfo> partitionInfos = getInfos()
            .values()
            .stream()
            .flatMap(Collection::stream)
            .collect(Collectors.toList());
    Cluster cluster = new Cluster(
            UUID.randomUUID().toString(),
            Collections.emptyList(),
            partitionInfos,
            Collections.emptySet(),
            Collections.emptySet()
    );
    return new MockProducer(cluster, true, null, new ByteBufferSerializer(), new ByteBufferSerializer());
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
private MockProducer createProducer() {
    Map<String, List<PartitionInfo>> infos = getInfos();
    Collection<PartitionInfo> partitionInfos = new ArrayList<>();
    for (List<PartitionInfo> infoEntry : infos.values()) {
        partitionInfos.addAll(infoEntry);
    }
    Cluster cluster = new Cluster(UUID.randomUUID().toString(), Collections.<Node>emptyList(), partitionInfos,
        Collections.<String>emptySet(), Collections.<String>emptySet());
    MockProducer producer = new MockProducer(cluster, true, null, new ByteBufferSerializer(), new ByteBufferSerializer());
    PRODUCERS.add(producer);
    return producer;
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@Test
public void shouldCloseAllTaskProducersOnCloseIfEosEnabled() {
    builder.addSource("source1", "someTopic");

    final MockClientSupplier clientSupplier = new MockClientSupplier(applicationId);
    final StreamThread thread = new StreamThread(
        builder,
        new StreamsConfig(configProps(true)),
        clientSupplier,
        applicationId,
        clientId,
        processId,
        metrics,
        mockTime,
        new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST),
        0);

    final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>();
    assignment.put(new TaskId(0, 0), Collections.singleton(new TopicPartition("someTopic", 0)));
    assignment.put(new TaskId(0, 1), Collections.singleton(new TopicPartition("someTopic", 1)));
    thread.setPartitionAssignor(new MockStreamsPartitionAssignor(assignment));

    thread.rebalanceListener.onPartitionsAssigned(Collections.singleton(new TopicPartition("someTopic", 0)));

    thread.close();
    thread.run();

    for (final StreamTask task : thread.tasks().values()) {
        assertTrue(((MockProducer) ((RecordCollectorImpl) task.recordCollector()).producer()).closed());
    }
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@Test
public void shouldCloseThreadProducerOnCloseIfEosDisabled() {
    builder.addSource("source1", "someTopic");

    final StreamThread thread = new StreamThread(
        builder,
        config,
        clientSupplier,
        applicationId,
        clientId,
        processId,
        metrics,
        mockTime,
        new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST),
        0);

    final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>();
    assignment.put(new TaskId(0, 0), Collections.singleton(new TopicPartition("someTopic", 0)));
    assignment.put(new TaskId(0, 1), Collections.singleton(new TopicPartition("someTopic", 1)));
    thread.setPartitionAssignor(new MockStreamsPartitionAssignor(assignment));

    thread.rebalanceListener.onPartitionsAssigned(Collections.singleton(new TopicPartition("someTopic", 0)));

    thread.close();
    thread.run();

    assertTrue(((MockProducer) thread.threadProducer).closed());
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@Test
public void testSpecificPartition() {

    final RecordCollectorImpl collector = new RecordCollectorImpl(
            new MockProducer<>(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer),
            "RecordCollectorTest-TestSpecificPartition");

    collector.send("topic1", "999", "0", 0, null, stringSerializer, stringSerializer);
    collector.send("topic1", "999", "0", 0, null, stringSerializer, stringSerializer);
    collector.send("topic1", "999", "0", 0, null, stringSerializer, stringSerializer);

    collector.send("topic1", "999", "0", 1, null, stringSerializer, stringSerializer);
    collector.send("topic1", "999", "0", 1, null, stringSerializer, stringSerializer);

    collector.send("topic1", "999", "0", 2, null, stringSerializer, stringSerializer);

    final Map<TopicPartition, Long> offsets = collector.offsets();

    assertEquals((Long) 2L, offsets.get(new TopicPartition("topic1", 0)));
    assertEquals((Long) 1L, offsets.get(new TopicPartition("topic1", 1)));
    assertEquals((Long) 0L, offsets.get(new TopicPartition("topic1", 2)));

    // ignore StreamPartitioner
    collector.send("topic1", "999", "0", 0, null, stringSerializer, stringSerializer);
    collector.send("topic1", "999", "0", 1, null, stringSerializer, stringSerializer);
    collector.send("topic1", "999", "0", 2, null, stringSerializer, stringSerializer);

    assertEquals((Long) 3L, offsets.get(new TopicPartition("topic1", 0)));
    assertEquals((Long) 2L, offsets.get(new TopicPartition("topic1", 1)));
    assertEquals((Long) 1L, offsets.get(new TopicPartition("topic1", 2)));
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Test
public void shouldRetryWhenTimeoutExceptionOccursOnSend() throws Exception {
    final AtomicInteger attempt = new AtomicInteger(0);
    final RecordCollectorImpl collector = new RecordCollectorImpl(
            new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
                @Override
                public synchronized Future<RecordMetadata> send(final ProducerRecord record, final Callback callback) {
                    if (attempt.getAndIncrement() == 0) {
                        throw new TimeoutException();
                    }
                    return super.send(record, callback);
                }
            },
            "test");

    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    final Long offset = collector.offsets().get(new TopicPartition("topic1", 0));
    assertEquals(Long.valueOf(0L), offset);
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Test(expected = StreamsException.class)
public void shouldThrowStreamsExceptionAfterMaxAttempts() throws Exception {
    final RecordCollector collector = new RecordCollectorImpl(
            new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
                @Override
                public synchronized Future<RecordMetadata> send(final ProducerRecord record, final Callback callback) {
                    throw new TimeoutException();
                }
            },
            "test");

    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);

}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Test(expected = StreamsException.class)
public void shouldThrowStreamsExceptionOnSubsequentCallIfASendFails() throws Exception {
    final RecordCollector collector = new RecordCollectorImpl(
            new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
                @Override
                public synchronized Future<RecordMetadata> send(final ProducerRecord record, final Callback callback) {
                    callback.onCompletion(null, new Exception());
                    return null;
                }
            },
            "test");
    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Test(expected = StreamsException.class)
public void shouldThrowStreamsExceptionOnFlushIfASendFailed() throws Exception {
    final RecordCollector collector = new RecordCollectorImpl(
            new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
                @Override
                public synchronized Future<RecordMetadata> send(final ProducerRecord record, final Callback callback) {
                    callback.onCompletion(null, new Exception());
                    return null;
                }
            },
            "test");
    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.flush();
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Test(expected = StreamsException.class)
public void shouldThrowStreamsExceptionOnCloseIfASendFailed() throws Exception {
    final RecordCollector collector = new RecordCollectorImpl(
            new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
                @Override
                public synchronized Future<RecordMetadata> send(final ProducerRecord record, final Callback callback) {
                    callback.onCompletion(null, new Exception());
                    return null;
                }
            },
            "test");
    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
    collector.close();
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Test(expected = StreamsException.class)
public void shouldThrowIfTopicIsUnknown() {
    final RecordCollector collector = new RecordCollectorImpl(
        new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
            @Override
            public List<PartitionInfo> partitionsFor(final String topic) {
                return Collections.EMPTY_LIST;
            }

        },
        "test");
    collector.send("topic1", "3", "0", null, stringSerializer, stringSerializer, streamPartitioner);
}
 

import org.apache.kafka.clients.producer.MockProducer; //导入依赖的package包/类
@Test
public void shouldInitAndBeginTransactionOnCreateIfEosEnabled() throws Exception {
    final MockProducer producer = new MockProducer();
    task = new StreamTask(taskId00, applicationId, partitions, topology, consumer, changelogReader,
        eosConfig, streamsMetrics, stateDirectory, null, time, producer);

    assertTrue(producer.transactionInitialized());
    assertTrue(producer.transactionInFlight());
}