Fix BigQuery Schema Provider test.

robertwb · robertwb · commit 63c16bbbe421 · 2023-09-08T08:33:30.000-07:00
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProviderTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProviderTest.java
@@ -48,9 +48,11 @@
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
 import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptors;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -211,7 +213,13 @@ public void testInputElementCount() throws Exception {
   public void testFailedRows() throws Exception {
     String tableSpec = "project:dataset.write_with_fail";
     BigQueryStorageWriteApiSchemaTransformConfiguration config =
-        BigQueryStorageWriteApiSchemaTransformConfiguration.builder().setTable(tableSpec).build();
+        BigQueryStorageWriteApiSchemaTransformConfiguration.builder()
+            .setTable(tableSpec)
+            .setErrorHandling(
+                BigQueryStorageWriteApiSchemaTransformConfiguration.ErrorHandling.builder()
+                    .setOutput("FailedRows")
+                    .build())
+            .build();
 
     String failValue = "fail_me";
 
@@ -234,7 +242,15 @@ public void testFailedRows() throws Exception {
     fakeDatasetService.setShouldFailRow(shouldFailRow);
 
     PCollectionRowTuple result = runWithConfig(config, totalRows);
-    PCollection<Row> failedRows = result.get("FailedRows");
+    PCollection<Row> failedRows =
+        result
+            .get("FailedRows")
+            .apply(
+                "ExtractFailedRows",
+                MapElements.into(TypeDescriptors.rows())
+                    .via((rowAndError) -> rowAndError.<Row>getValue("failed_row")))
+            .setRowSchema(SCHEMA);
+    ;
 
     PAssert.that(failedRows).containsInAnyOrder(expectedFailedRows);
     p.run().waitUntilFinish();
@@ -250,7 +266,13 @@ public void testFailedRows() throws Exception {
   public void testErrorCount() throws Exception {
     String tableSpec = "project:dataset.error_count";
     BigQueryStorageWriteApiSchemaTransformConfiguration config =
-        BigQueryStorageWriteApiSchemaTransformConfiguration.builder().setTable(tableSpec).build();
+        BigQueryStorageWriteApiSchemaTransformConfiguration.builder()
+            .setTable(tableSpec)
+            .setErrorHandling(
+                BigQueryStorageWriteApiSchemaTransformConfiguration.ErrorHandling.builder()
+                    .setOutput("FailedRows")
+                    .build())
+            .build();
 
     Function<TableRow, Boolean> shouldFailRow =
         (Function<TableRow, Boolean> & Serializable) tr -> tr.get("name").equals("a");
