htran1 commented on a change in pull request #2682: [GOBBLIN-820] Add keyed 
write capability to Kafka writer
URL: https://github.com/apache/incubator-gobblin/pull/2682#discussion_r301399447
 
 

 ##########
 File path: 
gobblin-modules/gobblin-kafka-09/src/test/java/org/apache/gobblin/kafka/writer/Kafka09DataWriterTest.java
 ##########
 @@ -180,5 +179,94 @@ public void testAvroSerialization()
   }
 
 
+  @Test
+  public void testKeyedAvroSerialization()
+      throws IOException, InterruptedException, SchemaRegistryException {
+    String topic = "testAvroSerialization09";
+    _kafkaTestHelper.provisionTopic(topic);
+    Properties props = new Properties();
+    props.setProperty(KafkaWriterConfigurationKeys.KAFKA_TOPIC, topic);
+    
props.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX + 
"bootstrap.servers",
+        "localhost:" + _kafkaTestHelper.getKafkaServerPort());
+    
props.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX + 
"value.serializer",
+        LiAvroSerializer.class.getName());
+    props.setProperty(KafkaWriterConfigurationKeys.WRITER_KAFKA_KEYED_CONFIG, 
"true");
+    String keyField = "field1";
+    
props.setProperty(KafkaWriterConfigurationKeys.WRITER_KAFKA_KEYFIELD_CONFIG, 
keyField);
+
+
+    // set up mock schema registry
+
+    props.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX
+            + KafkaSchemaRegistryConfigurationKeys.KAFKA_SCHEMA_REGISTRY_CLASS,
+        ConfigDrivenMd5SchemaRegistry.class.getCanonicalName());
+
+    Kafka09DataWriter<String, GenericRecord> kafka09DataWriter = new 
Kafka09DataWriter<>(props);
+    WriteCallback callback = mock(WriteCallback.class);
+
+    GenericRecord record = TestUtils.generateRandomAvroRecord();
+    try {
+      kafka09DataWriter.write(record, callback);
+    }
+    finally
+    {
+      kafka09DataWriter.close();
+    }
+
+    verify(callback, times(1)).onSuccess(isA(WriteResponse.class));
+    verify(callback, never()).onFailure(isA(Exception.class));
+    MessageAndMetadata<byte[], byte[]> value = 
_kafkaTestHelper.getIteratorForTopic(topic).next();
+    byte[] key = value.key();
+    byte[] message = value.message();
+    ConfigDrivenMd5SchemaRegistry schemaReg = new 
ConfigDrivenMd5SchemaRegistry(topic, record.getSchema());
+    LiAvroDeserializer deser = new LiAvroDeserializer(schemaReg);
+    GenericRecord receivedRecord = deser.deserialize(topic, message);
+    Assert.assertEquals(record.toString(), receivedRecord.toString());
+    Assert.assertEquals(new String(key), record.get(keyField));
+  }
+
+  @Test
+  public void testValueSerialization()
+      throws IOException, InterruptedException, SchemaRegistryException {
+    String topic = "testAvroSerialization09";
+    _kafkaTestHelper.provisionTopic(topic);
+    Properties props = new Properties();
+    props.setProperty(KafkaWriterConfigurationKeys.KAFKA_TOPIC, topic);
+    
props.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX + 
"bootstrap.servers",
+        "localhost:" + _kafkaTestHelper.getKafkaServerPort());
+    
props.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX + 
"value.serializer",
+    "org.apache.kafka.common.serialization.StringSerializer");
+    props.setProperty(KafkaWriterConfigurationKeys.WRITER_KAFKA_KEYED_CONFIG, 
"true");
+    String keyField = "field1";
+    
props.setProperty(KafkaWriterConfigurationKeys.WRITER_KAFKA_KEYFIELD_CONFIG, 
keyField);
+    
props.setProperty(KafkaWriterConfigurationKeys.WRITER_KAFKA_VALUEFIELD_CONFIG, 
keyField);
+
+
+    // set up mock schema registry
+
+    props.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX
+            + KafkaSchemaRegistryConfigurationKeys.KAFKA_SCHEMA_REGISTRY_CLASS,
+        ConfigDrivenMd5SchemaRegistry.class.getCanonicalName());
+
+    Kafka09DataWriter<String, GenericRecord> kafka09DataWriter = new 
Kafka09DataWriter<>(props);
+    WriteCallback callback = mock(WriteCallback.class);
+
+    GenericRecord record = TestUtils.generateRandomAvroRecord();
+    try {
+      kafka09DataWriter.write(record, callback);
+    }
+    finally
 
 Review comment:
   Format to have `finally` and the two braces on the same line?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


With regards,
Apache Git Services

Reply via email to