lukecwik commented on a change in pull request #9758: [BEAM-8374] Enable 
returning missing PublishResult fields in SnsIO.Write
URL: https://github.com/apache/beam/pull/9758#discussion_r391207323
 
 

 ##########
 File path: 
sdks/java/io/amazon-web-services/src/test/java/org/apache/beam/sdk/io/aws/sns/SnsIOTest.java
 ##########
 @@ -131,6 +139,52 @@ public void testRetries() throws Throwable {
     fail("Pipeline is expected to fail because we were unable to write to 
SNS.");
   }
 
+  @Test
+  public void testCustomCoder() throws Exception {
+    final PublishRequest request1 = createSampleMessage("my_first_message");
+
+    final TupleTag<PublishResult> results = new TupleTag<>();
+    final AmazonSNS amazonSnsSuccess = getAmazonSnsMockSuccess();
+    final MockCoder mockCoder = new MockCoder();
+
+    final PCollectionTuple snsWrites =
+        p.apply(Create.of(request1))
+            .apply(
+                SnsIO.write()
+                    .withTopicName(topicName)
+                    .withAWSClientsProvider(new Provider(amazonSnsSuccess))
+                    .withResultOutputTag(results)
+                    .withCoder(mockCoder));
+
+    final PCollection<Long> publishedResultsSize =
+        snsWrites
+            .get(results)
+            .apply(MapElements.into(TypeDescriptors.strings()).via(result -> 
result.getMessageId()))
+            .apply(Count.globally());
+    
PAssert.that(publishedResultsSize).containsInAnyOrder(ImmutableList.of(1L));
+    p.run().waitUntilFinish();
+    assertThat(mockCoder.captured).isNotNull();
+  }
+
+  // Hand-code mock because Mockito mocks cause NotSerializableException even 
with
+  // withSettings().serializable().
+  private static class MockCoder extends AtomicCoder<PublishResult> implements 
Serializable {
 
 Review comment:
   All Coders implement Serializable already.
   ```suggestion
     private static class MockCoder extends AtomicCoder<PublishResult> {
   ```

----------------------------------------------------------------
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