diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 3cc3c47acc54..801496e9d31e 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -4826,35 +4826,34 @@ void testLocationFastSQLQueryWithJobId() throws InterruptedException { /* TODO(prasmish): replicate the entire test case for executeSelect */ @Test void testFastSQLQueryMultiPage() throws InterruptedException { - String query = - "SELECT date, county, state_name, county_fips_code, confirmed_cases, deaths FROM " - + TABLE_ID_LARGE.getTable(); + String query = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 1005)) as num"; QueryJobConfiguration config = - QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build(); + QueryJobConfiguration.newBuilder(query) + .setDefaultDataset(DatasetId.of(DATASET)) + .setMaxResults(1000L) + .build(); + + java.util.function.Consumer assertResult = + r -> { + assertNotNull(r.getJobId()); + assertEquals(1, r.getSchema().getFields().size()); + assertEquals("num", r.getSchema().getFields().get(0).getName()); + assertEquals(LegacySQLTypeName.INTEGER, r.getSchema().getFields().get(0).getType()); + assertEquals(1005, r.getTotalRows()); + assertNotNull(r.getNextPage()); + assertNotNull(r.getNextPageToken()); + assertTrue(r.hasNextPage()); + }; + TableResult result = bigquery.query(config); - assertNotNull(result.getJobId()); - assertEquals(LARGE_TABLE_SCHEMA, result.getSchema()); - assertEquals(313348, result.getTotalRows()); - assertNotNull(result.getNextPage()); - assertNotNull(result.getNextPageToken()); - assertTrue(result.hasNextPage()); + assertResult.accept(result); TableResult result1 = bigquery.query(config); - assertNotNull(result1.getJobId()); - assertEquals(LARGE_TABLE_SCHEMA, result.getSchema()); - assertEquals(313348, result.getTotalRows()); - assertNotNull(result1.getNextPage()); - assertNotNull(result1.getNextPageToken()); - assertTrue(result1.hasNextPage()); + assertResult.accept(result1); - config.toBuilder().setQuery(query).build(); + config = config.toBuilder().setQuery(query).build(); TableResult result2 = bigquery.query(config); - assertNotNull(result2.getJobId()); - assertEquals(LARGE_TABLE_SCHEMA, result2.getSchema()); - assertEquals(313348, result2.getTotalRows()); - assertNotNull(result2.getNextPage()); - assertNotNull(result2.getNextPageToken()); - assertTrue(result2.hasNextPage()); + assertResult.accept(result2); } @Test