Fix Java GCP-IO Direct job

.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml

@@ -109,6 +109,7 @@ jobs:
             :sdks:java:io:google-cloud-platform:expansion-service:build \
             :sdks:java:io:google-cloud-platform:postCommit \
           arguments: |
+            --info \
             -PdisableSpotlessCheck=true \
             -PdisableCheckStyle=true \
             -PenableJacocoReport \

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java

@@ -131,7 +131,15 @@ public static Iterable<Object[]> data() {
   private static final int ORIGINAL_N = 60;
   // for dynamic destination test
   private static final int NUM_DESTINATIONS = 3;
-  private static final int TOTAL_NUM_STREAMS = 9;
+  private static final int TOTAL_NUM_STREAMS = 3;
+  // wait up to 60 seconds
+  private static final int SCHEMA_PROPAGATION_TIMEOUT_MS = 60000;
+  // interval between checks
+  private static final int SCHEMA_PROPAGATION_CHECK_INTERVAL_MS = 5000;
+  // wait for streams to recognize schema
+  private static final int STREAM_RECOGNITION_DELAY_MS = 15000;
+  // trigger for updating the schema when the row counter reaches this value
+  private static final int SCHEMA_UPDATE_TRIGGER = 1;
 
   private final Random randomGenerator = new Random();
 
@@ -218,17 +226,43 @@ public void setup() {
     public void processElement(ProcessContext c, @StateId(ROW_COUNTER) ValueState<Integer> counter)
         throws Exception {
       int current = firstNonNull(counter.read(), 0);
-      // We update schema early on to leave a healthy amount of time for StreamWriter to recognize
-      // it.
-      // We also update halfway through so that some writers are created *after* the schema update
-      if (current == TOTAL_NUM_STREAMS / 2) {
+      // We update schema early on to leave a healthy amount of time for the StreamWriter to recognize it,
+      // ensuring that subsequent writers are created with the updated schema.
+      if (current == SCHEMA_UPDATE_TRIGGER) {
         for (Map.Entry<String, String> entry : newSchemas.entrySet()) {
           bqClient.updateTableSchema(
               projectId,
               datasetId,
               entry.getKey(),
               BigQueryHelpers.fromJsonString(entry.getValue(), TableSchema.class));
         }
+
+        // check that schema update propagated fully
+        long startTime = System.currentTimeMillis();
+        long timeoutMillis = SCHEMA_PROPAGATION_TIMEOUT_MS;
+        boolean schemaPropagated = false;
+        while (System.currentTimeMillis() - startTime < timeoutMillis) {
+          schemaPropagated = true;
+          for (Map.Entry<String, String> entry : newSchemas.entrySet()) {
+            TableSchema currentSchema = bqClient.getTableResource(projectId, datasetId, entry.getKey()).getSchema();
+            TableSchema expectedSchema = BigQueryHelpers.fromJsonString(entry.getValue(), TableSchema.class);
+            if (currentSchema.getFields().size() != expectedSchema.getFields().size()) {
+              schemaPropagated = false;
+              break;
+            }
+          }
+          if (schemaPropagated) {
+            break;
+          }
+          Thread.sleep(SCHEMA_PROPAGATION_CHECK_INTERVAL_MS);
+        }
+        if (!schemaPropagated) {
+          LOG.warn("Schema update did not propagate fully within the timeout.");
+        } else {
+          LOG.info("Schema update propagated fully within the timeout - {}.", System.currentTimeMillis() - startTime);
+          // wait for streams to recognize the new schema
+          Thread.sleep(STREAM_RECOGNITION_DELAY_MS);
+        }
       }
 
       counter.write(++current);

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java

@@ -18,7 +18,6 @@
 package org.apache.beam.sdk.io.gcp.datastore;
 
 import static org.junit.Assert.assertThrows;
-import static org.mockito.Mockito.verify;
 
 import java.util.Map;
 import java.util.UUID;
@@ -50,7 +49,7 @@ public class RampupThrottlingFnTest {
   @Mock private Counter mockCounter;
   private final Sleeper mockSleeper =
       millis -> {
-        verify(mockCounter).inc(millis);
+        mockCounter.inc(millis);
         throw new RampupDelayException();
       };
   private DoFnTester<String, String> rampupThrottlingFnTester;

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java

@@ -70,6 +70,7 @@
 public class SpannerReadIT {
 
   private static final int MAX_DB_NAME_LENGTH = 30;
+  private static final int CLEANUP_PROPAGATION_DELAY_MS = 5000;
 
   @Rule public final transient TestPipeline p = TestPipeline.create();
   @Rule public transient ExpectedException thrown = ExpectedException.none();
@@ -285,6 +286,12 @@ private CloseTransactionFn(SpannerConfig spannerConfig) {
     public Transaction apply(Transaction tx) {
       BatchClient batchClient = SpannerAccessor.getOrCreate(spannerConfig).getBatchClient();
       batchClient.batchReadOnlyTransaction(tx.transactionId()).cleanup();
+      try {
+        // Wait for cleanup to propagate.
+        Thread.sleep(CLEANUP_PROPAGATION_DELAY_MS);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
       return tx;
     }
   }