DynamoDBIO { public static <T> Read<T> read() { return new AutoValue_DynamoDBIO_Read.Builder().build(); }  static Read<T> read(); static Write<T> write();  }
@Test public void testReadScanResult() { PCollection<List<Map<String, AttributeValue>>> actual = pipeline.apply( DynamoDBIO.<List<Map<String, AttributeValue>>>read() .withAwsClientsProvider( AwsClientsProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient())) .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> new ScanRequest(tableName).withTotalSegments(1)) .items()); PAssert.that(actual).containsInAnyOrder(expected); pipeline.run().waitUntilFinish(); }
@Test public void testMissingScanRequestFn() { thrown.expectMessage("withScanRequestFn() is required"); pipeline.apply( DynamoDBIO.read() .withAwsClientsProvider( AwsClientsProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))); try { pipeline.run().waitUntilFinish(); fail("withScanRequestFn() is required"); } catch (IllegalArgumentException ex) { assertEquals("withScanRequestFn() is required", ex.getMessage()); } }
@Test public void testMissingAwsClientsProvider() { thrown.expectMessage("withAwsClientsProvider() is required"); pipeline.apply( DynamoDBIO.read() .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> new ScanRequest(tableName).withTotalSegments(3))); try { pipeline.run().waitUntilFinish(); fail("withAwsClientsProvider() is required"); } catch (IllegalArgumentException ex) { assertEquals("withAwsClientsProvider() is required", ex.getMessage()); } }
@Test public void testMissingTotalSegments() { thrown.expectMessage("TotalSegments is required with withScanRequestFn()"); pipeline.apply( DynamoDBIO.read() .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> new ScanRequest(tableName)) .withAwsClientsProvider( AwsClientsProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))); try { pipeline.run().waitUntilFinish(); fail("TotalSegments is required with withScanRequestFn()"); } catch (IllegalArgumentException ex) { assertEquals("TotalSegments is required with withScanRequestFn()", ex.getMessage()); } }
@Test public void testNegativeTotalSegments() { thrown.expectMessage("TotalSegments is required with withScanRequestFn() and greater zero"); pipeline.apply( DynamoDBIO.read() .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> new ScanRequest(tableName).withTotalSegments(-1)) .withAwsClientsProvider( AwsClientsProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))); try { pipeline.run().waitUntilFinish(); fail("withTotalSegments() is expected and greater than zero"); } catch (IllegalArgumentException ex) { assertEquals( "TotalSegments is required with withScanRequestFn() and greater zero", ex.getMessage()); } }
@Test public void testReaderOneSegment() { List<Map<String, AttributeValue>> expected = DynamoDBIOTestHelper.generateTestData(tableName, numOfItems); PCollection<List<Map<String, AttributeValue>>> actual = pipeline.apply( DynamoDBIO.<List<Map<String, AttributeValue>>>read() .withDynamoDbClientProvider( DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient())) .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> ScanRequest.builder().tableName(tableName).totalSegments(1).build()) .items()); PAssert.that(actual).containsInAnyOrder(expected); pipeline.run().waitUntilFinish(); }
@Test public void testReaderThreeSegments() { TupleTag<List<Map<String, AttributeValue>>> outputTag = new TupleTag<>(); PCollectionTuple writeOutput = pipeline .apply( DynamoDBIO.<List<Map<String, AttributeValue>>>read() .withDynamoDbClientProvider( DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient())) .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> ScanRequest.builder().tableName(tableName).totalSegments(3).build()) .items()) .apply( ParDo.of( new DoFn< List<Map<String, AttributeValue>>, List<Map<String, AttributeValue>>>() { @ProcessElement public void processElement( @Element List<Map<String, AttributeValue>> input, OutputReceiver<List<Map<String, AttributeValue>>> out) { out.output(input); } }) .withOutputTags(outputTag, TupleTagList.empty())); final PCollection<Long> resultSetCount = writeOutput.get(outputTag).apply(Count.globally()); PAssert.that(resultSetCount).containsInAnyOrder(ImmutableList.of(3L)); pipeline.run().waitUntilFinish(); }
@Test public void testMissingScanRequestFn() { thrown.expectMessage("withScanRequestFn() is required"); pipeline.apply( DynamoDBIO.read() .withDynamoDbClientProvider( DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))); try { pipeline.run().waitUntilFinish(); fail("withScanRequestFn() is required"); } catch (IllegalArgumentException ex) { assertEquals("withScanRequestFn() is required", ex.getMessage()); } }
@Test public void testMissingDynamoDbClientProvider() { thrown.expectMessage("withDynamoDbClientProvider() is required"); pipeline.apply( DynamoDBIO.read() .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> ScanRequest.builder().tableName(tableName).totalSegments(3).build())); try { pipeline.run().waitUntilFinish(); fail("withDynamoDbClientProvider() is required"); } catch (IllegalArgumentException ex) { assertEquals("withDynamoDbClientProvider() is required", ex.getMessage()); } }
@Test public void testMissingTotalSegments() { thrown.expectMessage("TotalSegments is required with withScanRequestFn()"); pipeline.apply( DynamoDBIO.read() .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> ScanRequest.builder().tableName(tableName).build()) .withDynamoDbClientProvider( DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))); try { pipeline.run().waitUntilFinish(); fail("TotalSegments is required with withScanRequestFn()"); } catch (IllegalArgumentException ex) { assertEquals("TotalSegments is required with withScanRequestFn()", ex.getMessage()); } }
@Test public void testNegativeTotalSegments() { thrown.expectMessage("TotalSegments is required with withScanRequestFn() and greater zero"); pipeline.apply( DynamoDBIO.read() .withScanRequestFn( (SerializableFunction<Void, ScanRequest>) input -> ScanRequest.builder().tableName(tableName).totalSegments(-1).build()) .withDynamoDbClientProvider( DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))); try { pipeline.run().waitUntilFinish(); fail("withTotalSegments() is expected and greater than zero"); } catch (IllegalArgumentException ex) { assertEquals( "TotalSegments is required with withScanRequestFn() and greater zero", ex.getMessage()); } }