KafkaIO { public static <K, V> WriteRecords<K, V> writeRecords() { return new AutoValue_KafkaIO_WriteRecords.Builder<K, V>() .setProducerConfig(WriteRecords.DEFAULT_PRODUCER_PROPERTIES) .setEOS(false) .setNumShards(0) .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN) .build(); } private  KafkaIO(); static Read<byte[], byte[]> readBytes(); static Read<K, V> read(); static ReadSourceDescriptors<K, V> readSourceDescriptors(); static Write<K, V> write(); static WriteRecords<K, V> writeRecords();  }
@Test public void testRecordsSink() throws Exception { int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String topic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply(ParDo.of(new KV2ProducerRecord(topic))) .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) .apply( KafkaIO.<Integer, Long>writeRecords() .withBootstrapServers("none") .withTopic(topic) .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withInputTimestamp() .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); p.run(); completionThread.shutdown(); verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false, true); } }
@Test public void testSinkToMultipleTopics() throws Exception { int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String defaultTopic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply(ParDo.of(new KV2ProducerRecord(defaultTopic, false))) .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) .apply( KafkaIO.<Integer, Long>writeRecords() .withBootstrapServers("none") .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withInputTimestamp() .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); p.run(); completionThread.shutdown(); List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history(); for (int i = 0; i < numElements; i++) { ProducerRecord<Integer, Long> record = sent.get(i); if (i % 2 == 0) { assertEquals("test_2", record.topic()); } else { assertEquals("test_1", record.topic()); } assertEquals(i, record.key().intValue()); assertEquals(i, record.value().longValue()); assertEquals(i, record.timestamp().intValue()); assertEquals(0, record.headers().toArray().length); } } }
@Test public void testKafkaWriteHeaders() throws Exception { int numElements = 1; SimpleEntry<String, String> header = new SimpleEntry<>("header_key", "header_value"); try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String defaultTopic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply( ParDo.of( new KV2ProducerRecord(defaultTopic, true, System.currentTimeMillis(), header))) .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) .apply( KafkaIO.<Integer, Long>writeRecords() .withBootstrapServers("none") .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withInputTimestamp() .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); p.run(); completionThread.shutdown(); List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history(); for (int i = 0; i < numElements; i++) { ProducerRecord<Integer, Long> record = sent.get(i); Headers headers = record.headers(); assertNotNull(headers); Header[] headersArray = headers.toArray(); assertEquals(1, headersArray.length); assertEquals(header.getKey(), headersArray[0].key()); assertEquals( header.getValue(), new String(headersArray[0].value(), Charset.defaultCharset())); } } }
@Test public void testSinkProducerRecordsWithCustomTS() throws Exception { int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); final String defaultTopic = "test"; final Long ts = System.currentTimeMillis(); p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply(ParDo.of(new KV2ProducerRecord(defaultTopic, ts))) .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) .apply( KafkaIO.<Integer, Long>writeRecords() .withBootstrapServers("none") .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); p.run(); completionThread.shutdown(); List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history(); for (int i = 0; i < numElements; i++) { ProducerRecord<Integer, Long> record = sent.get(i); assertEquals(defaultTopic, record.topic()); assertEquals(i, record.key().intValue()); assertEquals(i, record.value().longValue()); assertEquals(ts, record.timestamp()); } } }