Expand test to include a proper table

we can connect to, and validate that
we can connect a BigQueryWriteClient to
the table.

Author: Eliassen-Steinar

modules/gcloud/src/test/java/org/testcontainers/containers/BigQueryEmulatorContainerTest.java

@@ -6,7 +6,16 @@
 import com.google.cloud.NoCredentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryOptions;
+import com.google.cloud.bigquery.DatasetId;
+import com.google.cloud.bigquery.DatasetInfo;
+import com.google.cloud.bigquery.Field;
 import com.google.cloud.bigquery.QueryJobConfiguration;
+import com.google.cloud.bigquery.Schema;
+import com.google.cloud.bigquery.StandardSQLTypeName;
+import com.google.cloud.bigquery.StandardTableDefinition;
+import com.google.cloud.bigquery.TableDefinition;
+import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.TableResult;
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
@@ -26,13 +35,71 @@
 
 class BigQueryEmulatorContainerTest {
 
+    private BigQuery getBigQuery(BigQueryEmulatorContainer container) {
+        String url = container.getEmulatorHttpEndpoint();
+        return BigQueryOptions
+            .newBuilder()
+            .setProjectId(container.getProjectId())
+            .setHost(url)
+            .setLocation(url)
+            .setCredentials(NoCredentials.getInstance())
+            .build().getService();
+    }
+
+    @Test
+    void test() throws Exception {
+        try (
+            // emulatorContainer {
+            BigQueryEmulatorContainer container = new BigQueryEmulatorContainer("ghcr.io/goccy/bigquery-emulator:0.4.3")
+            // }
+        ) {
+            container.start();
+
+            // bigQueryClient {
+            BigQuery bigQuery = getBigQuery(container);
+            // }
+
+            String fn =
+                "CREATE FUNCTION testr(arr ARRAY<STRUCT<name STRING, val INT64>>) AS ((SELECT SUM(IF(elem.name = \"foo\",elem.val,null)) FROM UNNEST(arr) AS elem))";
+
+            bigQuery.query(QueryJobConfiguration.newBuilder(fn).build());
+
+            String sql =
+                "SELECT testr([STRUCT<name STRING, val INT64>(\"foo\", 10), STRUCT<name STRING, val INT64>(\"bar\", 40), STRUCT<name STRING, val INT64>(\"foo\", 20)])";
+            TableResult result = bigQuery.query(QueryJobConfiguration.newBuilder(sql).build());
+            List<BigDecimal> values = result
+                .streamValues()
+                .map(fieldValues -> fieldValues.get(0).getNumericValue())
+                .collect(Collectors.toList());
+            assertThat(values).containsOnly(BigDecimal.valueOf(30));
+        }
+    }
+
     @Test
     public void testGrcp() throws IOException {
         // Shallow test, validate that connection can be set up, and attempt to create write stream fails.
         // BigQueryWriteSettings requires a HTTP/2 connection, not provided by the originally exposed endpoint. A "not found" exceptionm
         // indicates successful
         try (BigQueryEmulatorContainer container = new BigQueryEmulatorContainer("ghcr.io/goccy/bigquery-emulator:0.6.5")) {
             container.start();
+
+            // Create a table we can connect a write stream to
+            BigQuery bigQuery =  getBigQuery(container);
+            String tableName = "test-table";
+            String datasetName = "test-dataset";
+
+            bigQuery.create(DatasetInfo.of(DatasetId.of(container.getProjectId(), datasetName)));
+
+            Schema schema = Schema.of(
+                Field.of("name", StandardSQLTypeName.STRING)
+            );
+
+            TableId tableId = TableId.of(datasetName, tableName);
+            TableDefinition tableDefinition = StandardTableDefinition.of(schema);
+            TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();
+
+            bigQuery.create(tableInfo);
+
             BigQueryWriteSettings.Builder bigQueryWriteSettingsBuilder = BigQueryWriteSettings.newBuilder();
 
             bigQueryWriteSettingsBuilder.createWriteStreamSettings()
@@ -43,68 +110,24 @@ public void testGrcp() throws IOException {
                     .build());
 
             BigQueryWriteClient bigQueryWriteClient = BigQueryWriteClient.create(
-            bigQueryWriteSettingsBuilder.setTransportChannelProvider(FixedTransportChannelProvider.create(GrpcTransportChannel.create(
-                    ManagedChannelBuilder.forAddress(container.getHost(), container.getEmulatorGrpcPort()).usePlaintext().build())))
-                .setCredentialsProvider(NoCredentialsProvider.create())
-                .build()
+                bigQueryWriteSettingsBuilder.setTransportChannelProvider(FixedTransportChannelProvider.create(GrpcTransportChannel.create(
+                        ManagedChannelBuilder.forAddress(container.getHost(), container.getEmulatorGrpcPort()).usePlaintext().build())))
+                    .setCredentialsProvider(NoCredentialsProvider.create())
+                    .build()
             );
 
-            TableName parentTable = TableName.of(container.getProjectId(), "dataset", "table");
+            TableName parentTable = TableName.of(container.getProjectId(), datasetName, tableName);
             CreateWriteStreamRequest createWriteStreamRequest = CreateWriteStreamRequest.newBuilder()
                 .setParent(parentTable.toString())
                 .setWriteStream(WriteStream.newBuilder().setType(WriteStream.Type.PENDING))
                 .build();
 
-            String message = null;
-            try {
-                // This will fail, extract error message to check that it fails in a "we reached the backend" way to ensure that setup was correct
-                WriteStream writeStream =  bigQueryWriteClient.createWriteStream(createWriteStreamRequest);
-                // Example setting up StreamWriter. Note passing bigQueryWriteClient as parameter, this is needed to avoid using gcloud credentials:
-                /* StreamWriter writer = StreamWriter.newBuilder(writeStream.getName(), bigQueryWriteClient).setWriterSchema(schema).build(); */
-            } catch (RuntimeException e) {
-                message = e.getMessage();
-            }
-            assertThat(message).contains("dataset dataset is not found in project test-project");
+            // Validate that we can successfully create a write stream. This would not work with http endpoint
+            bigQueryWriteClient.createWriteStream(createWriteStreamRequest);
 
             bigQueryWriteClient.shutdown();
             bigQueryWriteClient.close();
         }
     }
 
-    @Test
-    void test() throws Exception {
-        try (
-            // emulatorContainer {
-            BigQueryEmulatorContainer container = new BigQueryEmulatorContainer("ghcr.io/goccy/bigquery-emulator:0.4.3")
-            // }
-        ) {
-            container.start();
-
-            // bigQueryClient {
-            String url = container.getEmulatorHttpEndpoint();
-            BigQueryOptions options = BigQueryOptions
-                .newBuilder()
-                .setProjectId(container.getProjectId())
-                .setHost(url)
-                .setLocation(url)
-                .setCredentials(NoCredentials.getInstance())
-                .build();
-            BigQuery bigQuery = options.getService();
-            // }
-
-            String fn =
-                "CREATE FUNCTION testr(arr ARRAY<STRUCT<name STRING, val INT64>>) AS ((SELECT SUM(IF(elem.name = \"foo\",elem.val,null)) FROM UNNEST(arr) AS elem))";
-
-            bigQuery.query(QueryJobConfiguration.newBuilder(fn).build());
-
-            String sql =
-                "SELECT testr([STRUCT<name STRING, val INT64>(\"foo\", 10), STRUCT<name STRING, val INT64>(\"bar\", 40), STRUCT<name STRING, val INT64>(\"foo\", 20)])";
-            TableResult result = bigQuery.query(QueryJobConfiguration.newBuilder(sql).build());
-            List<BigDecimal> values = result
-                .streamValues()
-                .map(fieldValues -> fieldValues.get(0).getNumericValue())
-                .collect(Collectors.toList());
-            assertThat(values).containsOnly(BigDecimal.valueOf(30));
-        }
-    }
 }