update test__run_using_base_image_with_mocks

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

@@ -7,6 +7,7 @@
 import pytest
 from pipelines.airbyte_ci.connectors.build_image.steps import build_customization, python_connectors
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
+from pipelines.consts import BUILD_PLATFORMS
 from pipelines.models.steps import StepStatus
 
 pytestmark = [
@@ -15,6 +16,10 @@
 
 
 class TestBuildConnectorImage:
+    @pytest.fixture
+    def all_platforms(self):
+        return BUILD_PLATFORMS
+
     @pytest.fixture
     def test_context(self, mocker):
         return mocker.Mock(secrets_to_mask=[])
@@ -95,18 +100,19 @@ def test_context_with_real_connector_without_base_image(self, connector_without_
         context.dagger_client = dagger_client
         return context
 
-    async def test__run_using_base_image_with_mocks(self, mocker, test_context_with_connector_with_base_image, current_platform):
+    async def test__run_using_base_image_with_mocks(self, mocker, test_context_with_connector_with_base_image, all_platforms):
         container_built_from_base = mocker.AsyncMock()
         mocker.patch.object(
             python_connectors.BuildConnectorImages, "_build_from_base_image", mocker.AsyncMock(return_value=container_built_from_base)
         )
         mocker.patch.object(python_connectors.BuildConnectorImages, "get_step_result", mocker.AsyncMock())
-        step = python_connectors.BuildConnectorImages(test_context_with_connector_with_base_image, current_platform)
+        step = python_connectors.BuildConnectorImages(test_context_with_connector_with_base_image, *all_platforms)
         step_result = await step._run()
-        step._build_from_base_image.assert_called_once()
-        container_built_from_base.with_exec.assert_called_once_with(["spec"])
+        assert step._build_from_base_image.call_count == len(all_platforms)
+        container_built_from_base.with_exec.assert_called_with(["spec"])
         assert step_result.status is StepStatus.SUCCESS
-        assert step_result.output_artifact[current_platform] == container_built_from_base
+        for platform in all_platforms:
+            assert step_result.output_artifact[platform] == container_built_from_base
 
     async def test_building_from_base_image_for_real(self, test_context_with_real_connector_using_base_image, current_platform):
         step = python_connectors.BuildConnectorImages(test_context_with_real_connector_using_base_image, current_platform)