JDBC Sources Integration Test migration follow-up: Switch additional …

…tests to use system stubs (#20040) * Migrating remaining tests * More tests * set env vars before querying them; use stream state everywhere? * fix strirct encrypt expected spec * runRead after workerConfigs exist * format Co-authored-by: Edward Gao <edward.gao@airbyte.io>
airbytehq · Dec 5, 2022 · 3b0512f · 3b0512f
1 parent 4f9b660
commit 3b0512f
Show file tree

Hide file tree

Showing 9 changed files with 442 additions and 23 deletions.
diff --git a/...rc/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java b/...rc/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java
@@ -125,6 +125,7 @@ public void setUpInternal() throws Exception {
     jobRoot = Files.createDirectories(Path.of(workspaceRoot.toString(), "job"));
     localRoot = Files.createTempDirectory(testDir, "output");
     environment = new TestDestinationEnv(localRoot);
+    setupEnvironment(environment);
     workerConfigs = new WorkerConfigs(new EnvConfigs());
     mConfigRepository = mock(ConfigRepository.class);
     processFactory = new DockerProcessFactory(
@@ -134,9 +135,15 @@ public void setUpInternal() throws Exception {
         localRoot.toString(),
         "host");
 
-    setupEnvironment(environment);
+    postSetup();
   }
 
+  /**
+   * Override this method if you want to do any per-test setup that depends on being able to e.g.
+   * {@link #runRead(ConfiguredAirbyteCatalog)}.
+   */
+  protected void postSetup() throws Exception {}
+
   @AfterEach
   public void tearDownInternal() throws Exception {
     tearDown(environment);

diff --git a/...airbyte/integration_tests/sources/AbstractPostgresSourceSSLCertificateAcceptanceTest.java b/...airbyte/integration_tests/sources/AbstractPostgresSourceSSLCertificateAcceptanceTest.java
@@ -9,6 +9,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
+import io.airbyte.commons.features.EnvVariableFeatureFlags;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.db.Database;
 import io.airbyte.db.PostgresUtils;
@@ -29,11 +30,19 @@
 import java.util.List;
 import org.jooq.DSLContext;
 import org.jooq.SQLDialect;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.utility.DockerImageName;
+import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
+import uk.org.webcompere.systemstubs.jupiter.SystemStub;
+import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension;
 
+@ExtendWith(SystemStubsExtension.class)
 public abstract class AbstractPostgresSourceSSLCertificateAcceptanceTest extends SourceAcceptanceTest {
 
+  @SystemStub
+  private EnvironmentVariables environmentVariables;
+
   private static final String STREAM_NAME = "public.id_and_name";
   private static final String STREAM_NAME2 = "public.starships";
   private static final String STREAM_NAME_MATERIALIZED_VIEW = "public.testview";
@@ -45,6 +54,8 @@ public abstract class AbstractPostgresSourceSSLCertificateAcceptanceTest extends
 
   @Override
   protected void setupEnvironment(final TestDestinationEnv environment) throws Exception {
+    environmentVariables.set(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, "true");
+
     container = new PostgreSQLContainer<>(DockerImageName.parse("postgres:bullseye")
         .asCompatibleSubstituteFor("postgres"));
     container.start();

diff --git a/...rations/io/airbyte/integration_tests/sources/AbstractSshPostgresSourceAcceptanceTest.java b/...rations/io/airbyte/integration_tests/sources/AbstractSshPostgresSourceAcceptanceTest.java
@@ -6,6 +6,7 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.Lists;
+import io.airbyte.commons.features.EnvVariableFeatureFlags;
 import io.airbyte.commons.functional.CheckedFunction;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.db.Database;
@@ -28,11 +29,19 @@
 import java.util.HashMap;
 import java.util.List;
 import org.jooq.SQLDialect;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.testcontainers.containers.Network;
 import org.testcontainers.containers.PostgreSQLContainer;
+import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
+import uk.org.webcompere.systemstubs.jupiter.SystemStub;
+import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension;
 
+@ExtendWith(SystemStubsExtension.class)
 public abstract class AbstractSshPostgresSourceAcceptanceTest extends SourceAcceptanceTest {
 
+  @SystemStub
+  private EnvironmentVariables environmentVariables;
+
   private static final String STREAM_NAME = "public.id_and_name";
   private static final String STREAM_NAME2 = "public.starships";
   private static final Network network = Network.newNetwork();
@@ -74,6 +83,7 @@ private static Database getDatabaseFromConfig(final JsonNode config) {
   // requiring data to already be in place.
   @Override
   protected void setupEnvironment(final TestDestinationEnv environment) throws Exception {
+    environmentVariables.set(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, "true");
     startTestContainers();
     config = bastion.getTunnelConfig(getTunnelMethod(), bastion.getBasicDbConfigBuider(db, List.of("public")));
     populateDatabaseTestData();

diff --git a/...tegrations/io/airbyte/integration_tests/sources/CdcWalLogsPostgresSourceDatatypeTest.java b/...tegrations/io/airbyte/integration_tests/sources/CdcWalLogsPostgresSourceDatatypeTest.java
@@ -42,7 +42,7 @@ protected List<AirbyteMessage> runRead(final ConfiguredAirbyteCatalog configured
   }
 
   @Override
-  protected void setupEnvironment(final TestDestinationEnv environment) throws Exception {
+  protected void postSetup() throws Exception {
     final Database database = setupDatabase();
     initTests();
     for (final TestDataHolder test : testDataHolders) {

diff --git a/...rbyte/integrations/io/airbyte/integration_tests/sources/PostgresSourceAcceptanceTest.java b/...rbyte/integrations/io/airbyte/integration_tests/sources/PostgresSourceAcceptanceTest.java
@@ -10,6 +10,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
+import io.airbyte.commons.features.EnvVariableFeatureFlags;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.db.Database;
 import io.airbyte.db.factory.DSLContextFactory;
@@ -35,10 +36,18 @@
 import org.jooq.DSLContext;
 import org.jooq.SQLDialect;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.testcontainers.containers.PostgreSQLContainer;
+import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
+import uk.org.webcompere.systemstubs.jupiter.SystemStub;
+import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension;
 
+@ExtendWith(SystemStubsExtension.class)
 public class PostgresSourceAcceptanceTest extends SourceAcceptanceTest {
 
+  @SystemStub
+  private EnvironmentVariables environmentVariables;
+
   private static final String STREAM_NAME = "public.id_and_name";
   private static final String STREAM_NAME2 = "public.starships";
   private static final String STREAM_NAME_MATERIALIZED_VIEW = "public.testview";
@@ -53,6 +62,8 @@ public class PostgresSourceAcceptanceTest extends SourceAcceptanceTest {
 
   @Override
   protected void setupEnvironment(final TestDestinationEnv environment) throws Exception {
+    environmentVariables.set(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, "true");
+
     container = new PostgreSQLContainer<>("postgres:13-alpine");
     container.start();
     String username = container.getUsername();

diff --git a/...tions/io/airbyte/integration_tests/sources/PostgresSourceStrictEncryptAcceptanceTest.java b/...tions/io/airbyte/integration_tests/sources/PostgresSourceStrictEncryptAcceptanceTest.java
@@ -9,6 +9,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
+import io.airbyte.commons.features.EnvVariableFeatureFlags;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.resources.MoreResources;
 import io.airbyte.db.Database;
@@ -30,18 +31,23 @@
 import java.util.HashMap;
 import org.jooq.DSLContext;
 import org.jooq.SQLDialect;
-import org.junitpioneer.jupiter.SetEnvironmentVariable;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.utility.DockerImageName;
+import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
+import uk.org.webcompere.systemstubs.jupiter.SystemStub;
+import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension;
 
 /**
  * This class is copied from source-postgres-strict-encrypt. The original file can be deleted
  * completely once the migration of multi-variant connector is done.
  */
-@SetEnvironmentVariable(key = "DEPLOYMENT_MODE",
-                        value = "CLOUD")
+@ExtendWith(SystemStubsExtension.class)
 public class PostgresSourceStrictEncryptAcceptanceTest extends SourceAcceptanceTest {
 
+  @SystemStub
+  private EnvironmentVariables environmentVariables;
+
   private static final String STREAM_NAME = "public.id_and_name";
   private static final String STREAM_NAME2 = "public.starships";
 
@@ -53,6 +59,8 @@ public class PostgresSourceStrictEncryptAcceptanceTest extends SourceAcceptanceT
 
   @Override
   protected void setupEnvironment(final TestDestinationEnv environment) throws Exception {
+    environmentVariables.set("DEPLOYMENT_MODE", "CLOUD");
+    environmentVariables.set(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, "true");
     container = new PostgreSQLContainer<>(DockerImageName.parse("postgres:bullseye")
         .asCompatibleSubstituteFor("postgres"));
     container.start();
@@ -109,7 +117,7 @@ protected String getImageName() {
 
   @Override
   protected ConnectorSpecification getSpec() throws Exception {
-    return SshHelpers.injectSshIntoSpec(Jsons.deserialize(MoreResources.readResource("expected_spec.json"), ConnectorSpecification.class));
+    return SshHelpers.injectSshIntoSpec(Jsons.deserialize(MoreResources.readResource("expected_strict_encrypt_spec.json"), ConnectorSpecification.class));
   }
 
   @Override