new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx);
@Test public void runQuery() throws Exception { SpannerIO.Read read = SpannerIO.read() .withProjectId("test") .withInstanceId("123") .withDatabaseId("aaa") .withQuery("SELECT * FROM users") .withServiceFactory(serviceFactory); List<Partition> fakePartitions = Arrays.asList(mock(Partition.class), mock(Partition.class), mock(Partition.class)); BatchTransactionId id = mock(BatchTransactionId.class); Transaction tx = Transaction.create(id); PCollectionView<Transaction> txView = pipeline.apply(Create.of(tx)).apply(View.<Transaction>asSingleton()); BatchSpannerRead.GeneratePartitionsFn fn = new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx); when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(id)).thenReturn(mockBatchTx); when(mockBatchTx.partitionQuery(any(PartitionOptions.class), any(Statement.class))) .thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray())); verify(serviceFactory.mockBatchClient()).batchReadOnlyTransaction(id); verify(mockBatchTx) .partitionQuery(any(PartitionOptions.class), eq(Statement.of("SELECT * " + "FROM users"))); }
new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx);
KV<ShardedKey<TableDestination>, List<String>>> tester = DoFnTester.of(writePartition); tester.setSideInput(tempFilePrefixView, GlobalWindow.INSTANCE, tempFilePrefix); tester.processElement(files);
tester.setSideInput(jobIdTokenView, GlobalWindow.INSTANCE, jobIdToken); tester.processElement(tempTablesElement); tester.finishBundle();
@Test public void fnWithSideInputExplicit() throws Exception { PCollection<Integer> pCollection = p.apply(Create.of(-2)); final PCollectionView<Integer> value = pCollection.apply(View.<Integer>asSingleton().withDefaultValue(0)); try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new SideInputDoFn(value))) { tester.setSideInput(value, GlobalWindow.INSTANCE, -2); tester.processElement(16); tester.processElement(32); tester.processElement(64); tester.processElement(128); tester.finishBundle(); assertThat(tester.peekOutputElements(), containsInAnyOrder(-2, -2, -2, -2)); } }