Parameterizing spanner Host in spanner tests (#2020)

* Parameterizing spanner Host parameter in spanner tests

* move testName cleanup logic to templateTestBase

* spotless

* Creating spannerTemplateITBase class for parameters for all SpannerITs

* Adding comments

* Adding a debug log

* Spotless fix

cicd/cmd/run-it-smoke-tests/main.go

@@ -66,7 +66,6 @@ func main() {
 		flags.ArtifactBucket(),
 		flags.StageBucket(),
 		flags.PrivateConnectivity(),
-		flags.SpannerHost(),
 		flags.FailureMode(),
 		flags.RetryFailures(),
 		flags.StaticOracleHost(),

cicd/cmd/run-it-tests/main.go

@@ -67,7 +67,6 @@ func main() {
 		flags.StageBucket(),
 		flags.HostIp(),
 		flags.PrivateConnectivity(),
-		flags.SpannerHost(),
 		flags.FailureMode(),
 		flags.RetryFailures(),
 		flags.StaticOracleHost(),

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/TemplateTestBase.java

@@ -90,11 +90,13 @@ public abstract class TemplateTestBase {
   public TestRule watcher =
       new TestWatcher() {
         protected void starting(Description description) {
-          LOG.info(
-              "Starting integration test {}.{}",
-              description.getClassName(),
-              description.getMethodName());
           testName = description.getMethodName();
+          // In case of parameterization the testName can contain subscript like testName[paramName]
+          // Converting testName from testName[paramName] to testNameParamName since it is used to
+          // create many resources and it cannot contain special characters.
+          testName = testName.replaceAll("\\[", "");
+          testName = testName.replaceAll("\\]", "");
+          LOG.info("Starting integration test {}.{}", description.getClassName(), testName);
         }
       };
 

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/SpannerResourceManager.java

@@ -81,7 +81,9 @@ public final class SpannerResourceManager implements ResourceManager {
   private static final Logger LOG = LoggerFactory.getLogger(SpannerResourceManager.class);
   private static final int MAX_BASE_ID_LENGTH = 30;
 
-  private static final String DEFAULT_SPANNER_HOST = "https://batch-spanner.googleapis.com";
+  public static final String DEFAULT_SPANNER_HOST = "https://batch-spanner.googleapis.com";
+  public static final String STAGING_SPANNER_HOST =
+      "https://staging-wrenchworks.sandbox.googleapis.com";
 
   // Retry settings for instance creation
   private static final int CREATE_MAX_RETRIES = 5;
@@ -232,7 +234,11 @@ private synchronized void maybeCreateDatabase() {
 
   private static <T> RetryPolicy<T> retryOnQuotaException() {
     return RetryPolicy.<T>builder()
-        .handleIf(exception -> ExceptionUtils.containsMessage(exception, "RESOURCE_EXHAUSTED"))
+        .handleIf(
+            exception -> {
+              LOG.warn("Error from spanner:", exception);
+              return ExceptionUtils.containsMessage(exception, "RESOURCE_EXHAUSTED");
+            })
         .withMaxRetries(CREATE_MAX_RETRIES)
         .withBackoff(CREATE_BACKOFF_DELAY, CREATE_BACKOFF_MAX_DELAY)
         .withJitter(CREATE_BACKOFF_JITTER)
@@ -597,15 +603,13 @@ public Builder setInstanceId(String instanceId) {
     }
 
     /**
-     * Looks at the system properties if there's a Spanner host override, uses it for Spanner API
-     * calls.
+     * Overrides spanner host, uses it for Spanner API calls.
      *
+     * @param spannerHost spanner host URL
      * @return this builder with host set.
      */
-    public Builder maybeUseCustomHost() {
-      if (System.getProperty("spannerHost") != null) {
-        this.host = System.getProperty("spannerHost");
-      }
+    public Builder useCustomHost(String spannerHost) {
+      this.host = spannerHost;
       return this;
     }
 

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/SpannerTemplateITBase.java

@@ -0,0 +1,42 @@
+package org.apache.beam.it.gcp.spanner;
+
+import java.util.Arrays;
+import java.util.Collection;
+import org.apache.beam.it.gcp.TemplateTestBase;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+/**
+ * Base class for all Spanner template integration tests. This class parameterizes the spannerHost
+ * for all subclasses. If the "spannerHost" system property is not set, the value of spannerHost
+ * will be set to STAGING_SPANNER_HOST and DEFAULT_SPANNER_HOST (as defined in {@link
+ * SpannerResourceManager}). All tests in the base class will be run twice: once with spannerHost
+ * set to STAGING_SPANNER_HOST and once with spannerHost set to DEFAULT_SPANNER_HOST. Otherwise, If
+ * the "spannerHost" system property is set, its value will be used to set spannerHost. All
+ * subclasses must use SpannerResourceManager.useCustomHost() and pass the spannerHost parameter to
+ * it.
+ */
+@RunWith(Parameterized.class)
+public abstract class SpannerTemplateITBase extends TemplateTestBase {
+
+  @Parameterized.Parameter(0)
+  public String spannerHost;
+
+  @Parameterized.Parameter(1)
+  public String spannerHostName;
+
+  // Because of parameterization, the test names will have subscripts. For example:
+  // testSpannerToGCSAvroBase[Staging]
+  @Parameters(name = "{1}")
+  public static Collection parameters() {
+    if (System.getProperty("spannerHost") != null) {
+      return Arrays.asList(new Object[][] {{System.getProperty("spannerHost"), "Custom"}});
+    }
+    return Arrays.asList(
+        new Object[][] {
+          {SpannerResourceManager.STAGING_SPANNER_HOST, "Staging"},
+          {SpannerResourceManager.DEFAULT_SPANNER_HOST, "Default"}
+        });
+  }
+}

v1/src/test/java/com/google/cloud/teleport/spanner/ExportPipelineIT.java

@@ -34,23 +34,23 @@
 import org.apache.beam.it.common.PipelineLauncher;
 import org.apache.beam.it.common.PipelineOperator;
 import org.apache.beam.it.common.utils.ResourceManagerUtils;
-import org.apache.beam.it.gcp.TemplateTestBase;
 import org.apache.beam.it.gcp.artifacts.Artifact;
 import org.apache.beam.it.gcp.artifacts.utils.AvroTestUtil;
 import org.apache.beam.it.gcp.spanner.SpannerResourceManager;
+import org.apache.beam.it.gcp.spanner.SpannerTemplateITBase;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.commons.lang3.RandomUtils;
 import org.junit.After;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
 
 /** Integration test for {@link ExportPipeline Spanner to GCS Avro} template. */
 @Category(TemplateIntegrationTest.class)
 @TemplateIntegrationTest(ExportPipeline.class)
-@RunWith(JUnit4.class)
-public class ExportPipelineIT extends TemplateTestBase {
+@RunWith(Parameterized.class)
+public class ExportPipelineIT extends SpannerTemplateITBase {
 
   private static final int MESSAGES_COUNT = 100;
 
@@ -154,16 +154,7 @@ public void testSpannerToGCSAvro() throws IOException {
     spannerResourceManager =
         SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.GOOGLE_STANDARD_SQL)
             .maybeUseStaticInstance()
-            .build();
-    testSpannerToGCSAvroBase(Function.identity());
-  }
-
-  @Test
-  public void testSpannerToGCSAvroStaging() throws IOException {
-    spannerResourceManager =
-        SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.GOOGLE_STANDARD_SQL)
-            .maybeUseStaticInstance()
-            .maybeUseCustomHost()
+            .useCustomHost(spannerHost)
             .build();
     testSpannerToGCSAvroBase(
         paramAdder ->
@@ -260,16 +251,7 @@ public void testPostgresSpannerToGCSAvro() throws IOException {
     spannerResourceManager =
         SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.POSTGRESQL)
             .maybeUseStaticInstance()
-            .build();
-    testPostgresSpannerToGCSAvroBase(Function.identity());
-  }
-
-  @Test
-  public void testPostgresSpannerToGCSAvroStaging() throws IOException {
-    spannerResourceManager =
-        SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.POSTGRESQL)
-            .maybeUseStaticInstance()
-            .maybeUseCustomHost()
+            .useCustomHost(spannerHost)
             .build();
     testPostgresSpannerToGCSAvroBase(
         paramAdder ->

v1/src/test/java/com/google/cloud/teleport/spanner/ImportPipelineIT.java

@@ -35,19 +35,19 @@
 import org.apache.beam.it.common.PipelineLauncher;
 import org.apache.beam.it.common.PipelineOperator;
 import org.apache.beam.it.common.utils.ResourceManagerUtils;
-import org.apache.beam.it.gcp.TemplateTestBase;
 import org.apache.beam.it.gcp.spanner.SpannerResourceManager;
+import org.apache.beam.it.gcp.spanner.SpannerTemplateITBase;
 import org.junit.After;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
 
 /** Integration test for {@link ImportPipeline} classic template. */
 @Category(TemplateIntegrationTest.class)
 @TemplateIntegrationTest(ImportPipeline.class)
-@RunWith(JUnit4.class)
-public class ImportPipelineIT extends TemplateTestBase {
+@RunWith(Parameterized.class)
+public class ImportPipelineIT extends SpannerTemplateITBase {
 
   private SpannerResourceManager spannerResourceManager;
 
@@ -127,16 +127,7 @@ public void testGoogleSqlImportPipeline() throws IOException {
     spannerResourceManager =
         SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.GOOGLE_STANDARD_SQL)
             .maybeUseStaticInstance()
-            .build();
-    testGoogleSqlImportPipelineBase(Function.identity());
-  }
-
-  @Test
-  public void testGoogleSqlImportPipelineStaging() throws IOException {
-    spannerResourceManager =
-        SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.GOOGLE_STANDARD_SQL)
-            .maybeUseStaticInstance()
-            .maybeUseCustomHost()
+            .useCustomHost(spannerHost)
             .build();
     testGoogleSqlImportPipelineBase(
         paramAdder ->
@@ -209,16 +200,7 @@ public void testPostgresImportPipeline() throws IOException {
     spannerResourceManager =
         SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.POSTGRESQL)
             .maybeUseStaticInstance()
-            .build();
-    testPostgresImportPipelineBase(Function.identity());
-  }
-
-  @Test
-  public void testPostgresImportPipelineStaging() throws IOException {
-    spannerResourceManager =
-        SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.POSTGRESQL)
-            .maybeUseStaticInstance()
-            .maybeUseCustomHost()
+            .useCustomHost(spannerHost)
             .build();
     testPostgresImportPipelineBase(
         paramAdder ->

v1/src/test/java/com/google/cloud/teleport/templates/SpannerToTextIT.java

@@ -32,21 +32,21 @@
 import org.apache.beam.it.common.PipelineLauncher;
 import org.apache.beam.it.common.PipelineOperator;
 import org.apache.beam.it.common.utils.ResourceManagerUtils;
-import org.apache.beam.it.gcp.TemplateTestBase;
 import org.apache.beam.it.gcp.artifacts.Artifact;
 import org.apache.beam.it.gcp.spanner.SpannerResourceManager;
+import org.apache.beam.it.gcp.spanner.SpannerTemplateITBase;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.After;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
 
 /** Integration test for {@link SpannerToText Spanner to GCS Text} template. */
 @Category(TemplateIntegrationTest.class)
 @TemplateIntegrationTest(SpannerToText.class)
-@RunWith(JUnit4.class)
-public class SpannerToTextIT extends TemplateTestBase {
+@RunWith(Parameterized.class)
+public class SpannerToTextIT extends SpannerTemplateITBase {
 
   private static final int MESSAGES_COUNT = 100;
 
@@ -62,16 +62,7 @@ public void testSpannerToGCSText() throws IOException {
     spannerResourceManager =
         SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.GOOGLE_STANDARD_SQL)
             .maybeUseStaticInstance()
-            .build();
-    testSpannerToGCSTextBase(Function.identity());
-  }
-
-  @Test
-  public void testSpannerToGCSTextStaging() throws IOException {
-    spannerResourceManager =
-        SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.GOOGLE_STANDARD_SQL)
-            .maybeUseStaticInstance()
-            .maybeUseCustomHost()
+            .useCustomHost(spannerHost)
             .build();
     testSpannerToGCSTextBase(
         paramAdder ->
@@ -145,16 +136,7 @@ public void testPostgresSpannerToGCSText() throws IOException {
     spannerResourceManager =
         SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.POSTGRESQL)
             .maybeUseStaticInstance()
-            .build();
-    testPostgresSpannerToGCSTextBase(Function.identity());
-  }
-
-  @Test
-  public void testPostgresSpannerToGCSTextStaging() throws IOException {
-    spannerResourceManager =
-        SpannerResourceManager.builder(testName, PROJECT, REGION, Dialect.POSTGRESQL)
-            .maybeUseStaticInstance()
-            .maybeUseCustomHost()
+            .useCustomHost(spannerHost)
             .build();
     testPostgresSpannerToGCSTextBase(
         paramAdder ->