Improving UT coverage for DatastreamToSpanner template (#3820)

Author: darshan-sj

v2/datastream-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/DataStreamToSpanner.java

@@ -82,9 +82,9 @@
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.TupleTagList;
-import org.joda.time.Duration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.threeten.bp.Duration;
 
 /**
  * This pipeline ingests DataStream data from GCS as events. The events are written to Cloud
@@ -616,6 +616,10 @@ static String getSourceType(Options options) {
       LOG.error("IOException Occurred: DataStreamClient failed initialization.");
       throw new IllegalArgumentException("Unable to initialize DatastreamClient: " + e);
     }
+    return getSourceTypeFromConfig(sourceConfig);
+  }
+
+  static String getSourceTypeFromConfig(SourceConfig sourceConfig) {
     if (sourceConfig.getMysqlSourceConfig() != null) {
       return DatastreamConstants.MYSQL_SOURCE_TYPE;
     } else if (sourceConfig.getOracleSourceConfig() != null) {
@@ -649,6 +653,10 @@ public static void main(String[] args) {
    * @return The result of the pipeline execution.
    */
   public static PipelineResult run(Options options) {
+    return buildPipeline(options).run();
+  }
+
+  static Pipeline buildPipeline(Options options) {
     long startTime = System.currentTimeMillis();
     /*
      * Stages:
@@ -682,13 +690,13 @@ public static PipelineResult run(Options options) {
             .withRpcPriority(ValueProvider.StaticValueProvider.of(options.getSpannerPriority()))
             .withCommitRetrySettings(
                 RetrySettings.newBuilder()
-                    .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(4))
-                    .setInitialRetryDelay(org.threeten.bp.Duration.ofMinutes(0))
+                    .setTotalTimeout(Duration.ofMinutes(4))
+                    .setInitialRetryDelay(Duration.ofMinutes(0))
                     .setRetryDelayMultiplier(1)
-                    .setMaxRetryDelay(org.threeten.bp.Duration.ofMinutes(0))
-                    .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(4))
+                    .setMaxRetryDelay(Duration.ofMinutes(0))
+                    .setInitialRpcTimeout(Duration.ofMinutes(4))
                     .setRpcTimeoutMultiplier(1)
-                    .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(4))
+                    .setMaxRpcTimeout(Duration.ofMinutes(4))
                     .setMaxAttempts(1)
                     .build());
     SpannerConfig shadowTableSpannerConfig = getShadowTableSpannerConfig(options);
@@ -776,7 +784,8 @@ public static PipelineResult run(Options options) {
                   .withFileReadConcurrency(options.getFileReadConcurrency())
                   .withoutDatastreamRecordsReshuffle()
                   .withDirectoryWatchDuration(
-                      Duration.standardMinutes(options.getDirectoryWatchDurationInMinutes()))
+                      org.joda.time.Duration.standardMinutes(
+                          options.getDirectoryWatchDurationInMinutes()))
                   .withDatastreamSourceType(options.getDatastreamSourceType()));
       int maxNumWorkers = options.getMaxNumWorkers() != 0 ? options.getMaxNumWorkers() : 1;
       jsonRecords =
@@ -855,7 +864,7 @@ public static PipelineResult run(Options options) {
     LOG.info("Filtered events directory: {}", filterEventsDirectory);
     transformedRecords
         .get(DatastreamToSpannerConstants.FILTERED_EVENT_TAG)
-        .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1))))
+        .apply(Window.into(FixedWindows.of(org.joda.time.Duration.standardMinutes(1))))
         .apply(
             "Write Filtered Events To GCS",
             TextIO.write().to(filterEventsDirectory).withSuffix(".json").withWindowedWrites());
@@ -928,8 +937,7 @@ public static PipelineResult run(Options options) {
                 .withTmpDirectory((options).getDeadLetterQueueDirectory() + "/tmp_severe/")
                 .setIncludePaneInfo(true)
                 .build());
-    // Execute the pipeline and return the result.
-    return pipeline.run();
+    return pipeline;
   }
 
   static SpannerConfig getShadowTableSpannerConfig(Options options) {
@@ -968,13 +976,13 @@ static SpannerConfig getShadowTableSpannerConfig(Options options) {
         .withRpcPriority(ValueProvider.StaticValueProvider.of(options.getSpannerPriority()))
         .withCommitRetrySettings(
             RetrySettings.newBuilder()
-                .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(4))
-                .setInitialRetryDelay(org.threeten.bp.Duration.ofMinutes(0))
+                .setTotalTimeout(Duration.ofMinutes(4))
+                .setInitialRetryDelay(Duration.ofMinutes(0))
                 .setRetryDelayMultiplier(1)
-                .setMaxRetryDelay(org.threeten.bp.Duration.ofMinutes(0))
-                .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(4))
+                .setMaxRetryDelay(Duration.ofMinutes(0))
+                .setInitialRpcTimeout(Duration.ofMinutes(4))
                 .setRpcTimeoutMultiplier(1)
-                .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(4))
+                .setMaxRpcTimeout(Duration.ofMinutes(4))
                 .setMaxAttempts(1)
                 .build());
   }

v2/datastream-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/spanner/ProcessInformationSchema.java

@@ -116,7 +116,10 @@ public ProcessInformationSchemaFn(
 
     @Setup
     public void setup() throws Exception {
-      spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);
+      // Use injected spannerAccessor if available (primarily for testing)
+      if (this.spannerAccessor == null) {
+        this.spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);
+      }
       DatabaseAdminClient databaseAdminClient = spannerAccessor.getDatabaseAdminClient();
       dialect =
           databaseAdminClient
@@ -129,7 +132,10 @@ public void setup() throws Exception {
         shadowTableSpannerAccessor = spannerAccessor;
         useSeparateShadowTableDb = false;
       } else {
-        shadowTableSpannerAccessor = SpannerAccessor.getOrCreate(shadowTableSpannerConfig);
+        // Use injected shadowTableSpannerAccessor if available (primarily for testing)
+        if (this.shadowTableSpannerAccessor == null) {
+          this.shadowTableSpannerAccessor = SpannerAccessor.getOrCreate(shadowTableSpannerConfig);
+        }
         useSeparateShadowTableDb = true;
       }
     }
@@ -203,10 +209,15 @@ Ddl getInformationSchemaAsDdl(SpannerAccessor accessor) {
       BatchClient batchClient = accessor.getBatchClient();
       BatchReadOnlyTransaction context =
           batchClient.batchReadOnlyTransaction(TimestampBound.strong());
-      InformationSchemaScanner scanner = new InformationSchemaScanner(context, dialect);
+      InformationSchemaScanner scanner = getInformationSchemaScanner(context, dialect);
       return scanner.scan();
     }
 
+    InformationSchemaScanner getInformationSchemaScanner(
+        BatchReadOnlyTransaction context, Dialect dialect) {
+      return new InformationSchemaScanner(context, dialect);
+    }
+
     void createShadowTablesInSpanner(Ddl mainDdl, Ddl shadowTableDdl) {
       List<String> dataTablesWithoutShadowTables =
           getDataTablesWithNoShadowTables(mainDdl, shadowTableDdl);