Search in sources :

Example 1 with CachingComposingSupplier

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();
}
Also used : AtomicInteger(java.util.concurrent.atomic.AtomicInteger) DatabaseClient(com.google.cloud.spanner.DatabaseClient) CachingComposingSupplier(org.springframework.cloud.gcp.data.spanner.core.admin.CachingComposingSupplier) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Test(org.junit.Test)

Aggregations

DatabaseClient (com.google.cloud.spanner.DatabaseClient)1 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)1 Test (org.junit.Test)1 CachingComposingSupplier (org.springframework.cloud.gcp.data.spanner.core.admin.CachingComposingSupplier)1