SnsIO { public static Write write() { return new AutoValue_SnsIO_Write.Builder().build(); }  static Write write();  }
@Test public void testDataWritesToSNS() { final PublishRequest request1 = createSampleMessage("my_first_message"); final PublishRequest request2 = createSampleMessage("my_second_message"); final TupleTag<PublishResult> results = new TupleTag<>(); final AmazonSNS amazonSnsSuccess = getAmazonSnsMockSuccess(); final PCollectionTuple snsWrites = p.apply(Create.of(request1, request2)) .apply( SnsIO.write() .withTopicName(topicName) .withRetryConfiguration( SnsIO.RetryConfiguration.create( 5, org.joda.time.Duration.standardMinutes(1))) .withAWSClientsProvider(new Provider(amazonSnsSuccess)) .withResultOutputTag(results)); final PCollection<Long> publishedResultsSize = snsWrites.get(results).apply(Count.globally()); PAssert.that(publishedResultsSize).containsInAnyOrder(ImmutableList.of(2L)); p.run().waitUntilFinish(); }
@Test public void testRetries() throws Throwable { thrown.expectMessage("Error writing to SNS"); final PublishRequest request1 = createSampleMessage("my message that will not be published"); final TupleTag<PublishResult> results = new TupleTag<>(); final AmazonSNS amazonSnsErrors = getAmazonSnsMockErrors(); p.apply(Create.of(request1)) .apply( SnsIO.write() .withTopicName(topicName) .withRetryConfiguration( SnsIO.RetryConfiguration.create(4, org.joda.time.Duration.standardSeconds(10))) .withAWSClientsProvider(new Provider(amazonSnsErrors)) .withResultOutputTag(results)); try { p.run(); } catch (final Pipeline.PipelineExecutionException e) { expectedLogs.verifyWarn(String.format(SnsIO.Write.SnsWriterFn.RETRY_ATTEMPT_LOG, 1)); expectedLogs.verifyWarn(String.format(SnsIO.Write.SnsWriterFn.RETRY_ATTEMPT_LOG, 2)); expectedLogs.verifyWarn(String.format(SnsIO.Write.SnsWriterFn.RETRY_ATTEMPT_LOG, 3)); throw e.getCause(); } 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(); }