test build customization

airbyte-ci/connectors/pipelines/tests/test_builds/dummy_build_customization.py

@@ -0,0 +1,35 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from dagger import Container
+
+
+async def pre_connector_install(base_image_container: Container) -> Container:
+    """This function will run before the connector installation.
+    It can mutate the base image container.
+
+    Args:
+        base_image_container (Container): The base image container to mutate.
+
+    Returns:
+        Container: The mutated base image container.
+    """
+    return await base_image_container.with_env_variable("MY_PRE_BUILD_ENV_VAR", "my_pre_build_env_var_value")
+
+
+async def post_connector_install(connector_container: Container) -> Container:
+    """This function will run after the connector installation during the build process.
+    It can mutate the connector container.
+
+    Args:
+        connector_container (Container): The connector container to mutate.
+
+    Returns:
+        Container: The mutated connector container.
+    """
+    return await connector_container.with_env_variable("MY_POST_BUILD_ENV_VAR", "my_post_build_env_var_value")

airbyte-ci/connectors/pipelines/tests/test_builds/test_python_connectors.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from pathlib import Path
+
 import pytest
 from pipelines.bases import StepStatus
 from pipelines.builds.python_connectors import BuildConnectorImage
@@ -31,9 +33,17 @@ def test_context_with_connector_without_base_image(self, test_context):
     def connector_with_base_image(self, all_connectors):
         for connector in all_connectors:
             if connector.metadata and connector.metadata.get("connectorBuildOptions", {}).get("baseImage"):
-                return connector
+                if not (connector.code_directory / "build_customization.py").exists():
+                    return connector
         pytest.skip("No connector with a connectorBuildOptions.baseImage metadata found")
 
+    @pytest.fixture
+    def connector_with_base_image_with_build_customization(self, connector_with_base_image):
+        dummy_build_customization = (Path(__file__).parent / "dummy_build_customization.py").read_text()
+        (connector_with_base_image.code_directory / "build_customization.py").write_text(dummy_build_customization)
+        yield connector_with_base_image
+        (connector_with_base_image.code_directory / "build_customization.py").unlink()
+
     @pytest.fixture
     def test_context_with_real_connector_using_base_image(self, connector_with_base_image, dagger_client):
         context = ConnectorContext(
@@ -48,6 +58,22 @@ def test_context_with_real_connector_using_base_image(self, connector_with_base_
         context.dagger_client = dagger_client
         return context
 
+    @pytest.fixture
+    def test_context_with_real_connector_using_base_image_with_build_customization(
+        self, connector_with_base_image_with_build_customization, dagger_client
+    ):
+        context = ConnectorContext(
+            pipeline_name="test build",
+            connector=connector_with_base_image_with_build_customization,
+            git_branch="test",
+            git_revision="test",
+            report_output_prefix="test",
+            is_local=True,
+            use_remote_secrets=True,
+        )
+        context.dagger_client = dagger_client
+        return context
+
     @pytest.fixture
     def connector_without_base_image(self, all_connectors):
         for connector in all_connectors:
@@ -85,9 +111,9 @@ async def test_building_from_base_image_for_real(self, test_context_with_real_co
         step_result = await step._run()
         step_result.status is StepStatus.SUCCESS
         built_container = step_result.output_artifact
-        assert await built_container.env_variable("AIRBYTE_ENTRYPOINT") == " ".join(step.DEFAULT_ENTRYPOINT)
+        assert await built_container.env_variable("AIRBYTE_ENTRYPOINT") == " ".join(step.get_entrypoint("main.py"))
         assert await built_container.workdir() == step.PATH_TO_INTEGRATION_CODE
-        assert await built_container.entrypoint() == step.DEFAULT_ENTRYPOINT
+        assert await built_container.entrypoint() == step.get_entrypoint("main.py")
         assert (
             await built_container.label("io.airbyte.version")
             == test_context_with_real_connector_using_base_image.connector.metadata["dockerImageTag"]
@@ -97,6 +123,16 @@ async def test_building_from_base_image_for_real(self, test_context_with_real_co
             == test_context_with_real_connector_using_base_image.connector.metadata["dockerRepository"]
         )
 
+    async def test_building_from_base_image_with_customization_for_real(
+        self, test_context_with_real_connector_using_base_image_with_build_customization, current_platform
+    ):
+        step = BuildConnectorImage(test_context_with_real_connector_using_base_image_with_build_customization, current_platform)
+        step_result = await step._run()
+        step_result.status is StepStatus.SUCCESS
+        built_container = step_result.output_artifact
+        assert await built_container.env_variable("MY_PRE_BUILD_ENV_VAR") == "my_pre_build_env_var_value"
+        assert await built_container.env_variable("MY_POST_BUILD_ENV_VAR") == "my_post_build_env_var_value"
+
     async def test__run_using_base_dockerfile_with_mocks(self, mocker, test_context_with_connector_without_base_image, current_platform):
         container_built_from_dockerfile = mocker.Mock()
         mocker.patch.object(BuildConnectorImage, "_build_from_dockerfile", mocker.AsyncMock(return_value=container_built_from_dockerfile))