@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")));
}