Skip to content

Commit 316dbc8

Browse files
committed
adding integration test
1 parent d391d50 commit 316dbc8

File tree

2 files changed

+55
-0
lines changed

2 files changed

+55
-0
lines changed

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,4 +196,8 @@ public void shutdownNow() {
196196
public boolean awaitTermination(long duration, TimeUnit unit) throws InterruptedException {
197197
return stub.awaitTermination(duration, unit);
198198
}
199+
200+
public BigQueryReadStubSettings getStubSettings() {
201+
return stubSettings;
202+
}
199203
}

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@
2121
import static org.junit.Assert.assertEquals;
2222
import static org.junit.Assert.assertNotNull;
2323
import static org.junit.Assert.assertNull;
24+
import static org.junit.Assert.assertTrue;
2425

26+
import com.google.api.gax.core.InstantiatingExecutorProvider;
2527
import com.google.api.gax.rpc.ServerStream;
2628
import com.google.cloud.RetryOption;
2729
import com.google.cloud.ServiceOptions;
@@ -44,6 +46,7 @@
4446
import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
4547
import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
4648
import com.google.cloud.bigquery.storage.v1.ReadSession;
49+
import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings;
4750
import com.google.cloud.bigquery.storage.v1.ReadSession.TableModifiers;
4851
import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions;
4952
import com.google.cloud.bigquery.storage.v1.ReadStream;
@@ -806,6 +809,54 @@ public void testStructAndArraySqlTypes() throws InterruptedException, IOExceptio
806809
assertEquals(rowAssertMessage, new Utf8("abc"), structRecord.get("str_field"));
807810
}
808811

812+
@Test
813+
public void testSimpleReadWithBackgroundExecutorProvider() throws IOException {
814+
BigQueryReadSettings bigQueryReadSettings = BigQueryReadSettings
815+
.newBuilder()
816+
.setBackgroundExecutorProvider(
817+
InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(14).build())
818+
.build();
819+
// Overriding the default client
820+
client = BigQueryReadClient.create(bigQueryReadSettings);
821+
assertTrue(
822+
client.getStub().getStubSettings().getBackgroundExecutorProvider() instanceof InstantiatingExecutorProvider);
823+
assertEquals(
824+
14,
825+
((InstantiatingExecutorProvider) client.getStub().getStubSettings().getBackgroundExecutorProvider())
826+
.getExecutorThreadCount());
827+
String table =
828+
BigQueryResource.FormatTableResource(
829+
/* projectId = */ "bigquery-public-data",
830+
/* datasetId = */ "samples",
831+
/* tableId = */ "shakespeare");
832+
833+
ReadSession session =
834+
client.createReadSession(
835+
/* parent = */ parentProjectId,
836+
/* readSession = */ ReadSession.newBuilder()
837+
.setTable(table)
838+
.setDataFormat(DataFormat.AVRO)
839+
.build(),
840+
/* maxStreamCount = */ 1);
841+
assertEquals(
842+
String.format(
843+
"Did not receive expected number of streams for table '%s' CreateReadSession response:%n%s",
844+
table, session.toString()),
845+
1,
846+
session.getStreamsCount());
847+
848+
ReadRowsRequest readRowsRequest =
849+
ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build();
850+
851+
long rowCount = 0;
852+
ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
853+
for (ReadRowsResponse response : stream) {
854+
rowCount += response.getRowCount();
855+
}
856+
857+
assertEquals(164_656, rowCount);
858+
}
859+
809860
/**
810861
* Reads to the specified row offset within the stream. If the stream does not have the desired
811862
* rows to read, it will read all of them.

0 commit comments

Comments
 (0)