use of org.springframework.cloud.gcp.data.spanner.core.admin.CachingComposingSupplier in project spring-cloud-gcp by spring-cloud.
the class SpannerTemplateTests method multipleDatabaseClientTest.
@Test
public void multipleDatabaseClientTest() {
DatabaseClient databaseClient1 = mock(DatabaseClient.class);
DatabaseClient databaseClient2 = mock(DatabaseClient.class);
when(databaseClient1.singleUse()).thenReturn(this.readContext);
when(databaseClient2.singleUse()).thenReturn(this.readContext);
AtomicInteger currentClient = new AtomicInteger(1);
Supplier<Integer> regionProvider = currentClient::getAndIncrement;
// this client selector will alternate between the two clients
Supplier<DatabaseClient> clientProvider = new CachingComposingSupplier<>(regionProvider, u -> u % 2 == 1 ? databaseClient1 : databaseClient2);
SpannerTemplate template = new SpannerTemplate(clientProvider, this.mappingContext, this.objectMapper, this.mutationFactory, this.schemaUtils);
// this first read should use the first client
template.read(TestEntity.class, Key.of("key"));
verify(databaseClient1, times(1)).singleUse();
verify(databaseClient2, times(0)).singleUse();
// this second read should use the second client
template.read(TestEntity.class, Key.of("key"));
verify(databaseClient1, times(1)).singleUse();
verify(databaseClient2, times(1)).singleUse();
// this third read should use the first client again
template.read(TestEntity.class, Key.of("key"));
verify(databaseClient1, times(2)).singleUse();
verify(databaseClient2, times(1)).singleUse();
}
Aggregations