Allow us to pass a map of environment variables within our BasicFunct…

…ionalityIntegrationTest class (#50427)

Co-authored-by: Octavia Squidington III <octavia-squidington-iii@users.noreply.github.com>

airbyte-cdk/bulk/core/load/src/integrationTest/kotlin/io/airbyte/cdk/load/mock_integration_test/MockBasicFunctionalityIntegrationTest.kt

@@ -27,6 +27,7 @@ class MockBasicFunctionalityIntegrationTest :
         preserveUndeclaredFields = true,
         commitDataIncrementally = false,
         allTypesBehavior = Untyped,
+        envVars = emptyMap(),
         supportFileTransfer = false,
     ) {
     @Test

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/IntegrationTest.kt

@@ -129,13 +129,15 @@ abstract class IntegrationTest(
         messages: List<InputMessage>,
         streamStatus: AirbyteStreamStatus? = AirbyteStreamStatus.COMPLETE,
         useFileTransfer: Boolean = false,
+        envVars: Map<String, String> = emptyMap(),
     ): List<AirbyteMessage> =
         runSync(
             configContents,
             DestinationCatalog(listOf(stream)),
             messages,
             streamStatus,
-            useFileTransfer
+            useFileTransfer,
+            envVars
         )
 
     /**
@@ -170,13 +172,15 @@ abstract class IntegrationTest(
          */
         streamStatus: AirbyteStreamStatus? = AirbyteStreamStatus.COMPLETE,
         useFileTransfer: Boolean = false,
+        envVars: Map<String, String> = emptyMap(),
     ): List<AirbyteMessage> {
         val destination =
             destinationProcessFactory.createDestinationProcess(
                 "write",
                 configContents,
                 catalog.asProtocolObject(),
                 useFileTransfer = useFileTransfer,
+                envVars = envVars,
             )
         return runBlocking(Dispatchers.IO) {
             launch { destination.run() }
@@ -212,13 +216,15 @@ abstract class IntegrationTest(
         inputStateMessage: StreamCheckpoint,
         allowGracefulShutdown: Boolean,
         useFileTransfer: Boolean = false,
+        envVars: Map<String, String> = emptyMap(),
     ): AirbyteStateMessage {
         val destination =
             destinationProcessFactory.createDestinationProcess(
                 "write",
                 configContents,
                 DestinationCatalog(listOf(stream)).asProtocolObject(),
                 useFileTransfer,
+                envVars
             )
         return runBlocking(Dispatchers.IO) {
             launch {

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/destination_process/DestinationProcess.kt

@@ -63,6 +63,7 @@ abstract class DestinationProcessFactory {
         configContents: String? = null,
         catalog: ConfiguredAirbyteCatalog? = null,
         useFileTransfer: Boolean = false,
+        envVars: Map<String, String> = emptyMap(),
         vararg featureFlags: FeatureFlag,
     ): DestinationProcess
 

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/destination_process/DockerizedDestination.kt

@@ -41,6 +41,7 @@ class DockerizedDestination(
     catalog: ConfiguredAirbyteCatalog?,
     private val testName: String,
     useFileTransfer: Boolean,
+    envVars: Map<String, String>,
     vararg featureFlags: FeatureFlag,
 ) : DestinationProcess {
     private val process: Process
@@ -94,6 +95,8 @@ class DockerizedDestination(
         val containerName = "$shortImageName-$command-$randomSuffix"
         logger.info { "Creating docker container $containerName" }
         logger.info { "File transfer ${if (useFileTransfer) "is " else "isn't"} enabled" }
+        val additionalEnvEntries = envVars.flatMap { (key, value) -> listOf("-e", "$key=$value") }
+        logger.info { "Env vars: $envVars loaded" }
         val cmd: MutableList<String> =
             (listOf(
                     "docker",
@@ -120,6 +123,7 @@ class DockerizedDestination(
                     "-e",
                     "USE_FILE_TRANSFER=$useFileTransfer",
                 ) +
+                    additionalEnvEntries +
                     featureFlags.flatMap { listOf("-e", it.envVarBindingDeclaration) } +
                     listOf(
                         // Yes, we hardcode the job ID to 0.
@@ -275,6 +279,7 @@ class DockerizedDestinationFactory(
         configContents: String?,
         catalog: ConfiguredAirbyteCatalog?,
         useFileTransfer: Boolean,
+        envVars: Map<String, String>,
         vararg featureFlags: FeatureFlag,
     ): DestinationProcess {
         return DockerizedDestination(
@@ -284,6 +289,7 @@ class DockerizedDestinationFactory(
             catalog,
             testName,
             useFileTransfer,
+            envVars,
             *featureFlags,
         )
     }

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/destination_process/NonDockerizedDestination.kt

@@ -12,6 +12,7 @@ import io.airbyte.cdk.load.test.util.IntegrationTest
 import io.airbyte.cdk.load.util.serializeToString
 import io.airbyte.protocol.models.v0.AirbyteMessage
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog
+import io.github.oshai.kotlinlogging.KotlinLogging
 import java.io.File
 import java.io.PipedInputStream
 import java.io.PipedOutputStream
@@ -23,11 +24,14 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import org.junit.jupiter.api.Assertions.assertFalse
 
+private val logger = KotlinLogging.logger {}
+
 class NonDockerizedDestination(
     command: String,
     configContents: String?,
     catalog: ConfiguredAirbyteCatalog?,
     useFileTransfer: Boolean,
+    envVars: Map<String, String>,
     vararg featureFlags: FeatureFlag,
 ) : DestinationProcess {
     private val destinationStdinPipe: PrintWriter
@@ -41,6 +45,9 @@ class NonDockerizedDestination(
     private val file = File("/tmp/test_file")
 
     init {
+        envVars.forEach { (key, value) -> IntegrationTest.nonDockerMockEnvVars.set(key, value) }
+        logger.info { "Env vars: $envVars loaded" }
+
         if (useFileTransfer) {
             IntegrationTest.nonDockerMockEnvVars.set("USE_FILE_TRANSFER", "true")
             val fileContentStr = "123"
@@ -118,6 +125,7 @@ class NonDockerizedDestinationFactory : DestinationProcessFactory() {
         configContents: String?,
         catalog: ConfiguredAirbyteCatalog?,
         useFileTransfer: Boolean,
+        envVars: Map<String, String>,
         vararg featureFlags: FeatureFlag,
     ): DestinationProcess {
         // TODO pass test name into the destination process
@@ -126,6 +134,7 @@ class NonDockerizedDestinationFactory : DestinationProcessFactory() {
             configContents,
             catalog,
             useFileTransfer,
+            envVars,
             *featureFlags
         )
     }

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/write/BasicFunctionalityIntegrationTest.kt

@@ -116,6 +116,7 @@ abstract class BasicFunctionalityIntegrationTest(
     val promoteUnionToObject: Boolean,
     val preserveUndeclaredFields: Boolean,
     val supportFileTransfer: Boolean,
+    val envVars: Map<String, String>,
     /**
      * Whether the destination commits new data when it receives a non-`COMPLETE` stream status. For
      * example:
@@ -170,7 +171,8 @@ abstract class BasicFunctionalityIntegrationTest(
                         blob = """{"foo": "bar"}""",
                         sourceRecordCount = 1,
                     )
-                )
+                ),
+                envVars = envVars,
             )
 
         val stateMessages = messages.filter { it.type == AirbyteMessage.Type.STATE }

airbyte-integrations/connectors/destination-dev-null/src/test-integration/kotlin/io/airbyte/integrations/destination/dev_null/DevNullBasicFunctionalityIntegrationTest.kt

@@ -26,6 +26,7 @@ class DevNullBasicFunctionalityIntegrationTest :
         preserveUndeclaredFields = false,
         commitDataIncrementally = false,
         allTypesBehavior = Untyped,
+        envVars = emptyMap(),
         supportFileTransfer = false,
     ) {
     @Test

airbyte-integrations/connectors/destination-iceberg-v2/src/test-integration/kotlin/io/airbyte/integrations/destination/iceberg/v2/IcebergV2WriteTest.kt

@@ -23,6 +23,7 @@ import org.junit.jupiter.api.Test
 abstract class IcebergV2WriteTest(
     configContents: String,
     destinationCleaner: DestinationCleaner,
+    envVars: Map<String, String> = emptyMap(),
 ) :
     BasicFunctionalityIntegrationTest(
         configContents,
@@ -37,6 +38,7 @@ abstract class IcebergV2WriteTest(
         preserveUndeclaredFields = false,
         commitDataIncrementally = false,
         supportFileTransfer = false,
+        envVars = envVars,
         allTypesBehavior = StronglyTyped(),
         nullEqualsUnset = true,
     ) {
@@ -80,11 +82,16 @@ class IcebergGlueWriteTest :
             IcebergV2TestUtil.getCatalog(
                 IcebergV2TestUtil.parseConfig(IcebergV2TestUtil.GLUE_CONFIG_PATH)
             )
-        ),
+        )
     ) {
     @Test
     @Disabled("dest iceberge-v2 doesn't support unknown types")
     override fun testUnknownTypes() {}
+
+    @Test
+    override fun testBasicWrite() {
+        super.testBasicWrite()
+    }
 }
 
 @Disabled(
@@ -96,6 +103,7 @@ class IcebergNessieMinioWriteTest :
         // we're writing to ephemeral testcontainers, so no need to clean up after ourselves
         NoopDestinationCleaner
     ) {
+
     companion object {
         private fun getToken(): String {
             val client = OkHttpClient()

airbyte-integrations/connectors/destination-s3-v2/src/test-integration/kotlin/io/airbyte/integrations/destination/s3_v2/S3V2WriteTest.kt

@@ -26,6 +26,7 @@ abstract class S3V2WriteTest(
     allTypesBehavior: AllTypesBehavior,
     nullEqualsUnset: Boolean = false,
     failOnUnknownTypes: Boolean = false,
+    envVars: Map<String, String> = emptyMap(),
 ) :
     BasicFunctionalityIntegrationTest(
         S3V2TestUtils.getConfig(path),
@@ -42,6 +43,7 @@ abstract class S3V2WriteTest(
         allTypesBehavior = allTypesBehavior,
         nullEqualsUnset = nullEqualsUnset,
         supportFileTransfer = true,
+        envVars = envVars,
         failOnUnknownTypes = failOnUnknownTypes,
     ) {
     @Disabled("Irrelevant for file destinations")