Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,7 @@ def _invoke_docker_build_and_push(self, container_image_name):
build = cloud_build_types.Build()
if self._cloud_build_machine_type:
build.options = cloud_build_types.BuildOptions()
build.options.machineType = self._cloud_build_machine_type
build.options.machine_type = self._cloud_build_machine_type
build.steps = []
step = cloud_build_types.BuildStep()
step.name = 'quay.io/buildah/stable:latest'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,18 @@

import gc
import logging
import tempfile
import unittest
import unittest.mock

from apache_beam.options import pipeline_options
from apache_beam.runners.portability import sdk_container_builder

try:
from google.cloud.devtools import cloudbuild_v1
except ImportError:
cloudbuild_v1 = None


class SdkContainerBuilderTest(unittest.TestCase):
def tearDown(self):
Expand Down Expand Up @@ -120,6 +126,41 @@ def test_container_image_name_uses_forward_slashes(self):
f' got: {container_image_name}'
)

@unittest.skipIf(
cloudbuild_v1 is None, 'GCP Cloud Build dependencies are not installed')
def test_cloud_builder_assign_machine_type(self):
"""Verify Cloud Build machine_type is correctly set on BuildOptions."""
from google.cloud.devtools.cloudbuild_v1 import types as cloud_build_types

with tempfile.TemporaryDirectory() as temp_dir:
options = pipeline_options.PipelineOptions([
'--project=test-project',
'--temp_location=gs://test-bucket/temp',
'--cloud_build_machine_type=n1-highcpu-8',
'--no_auth', # Avoid fetching credential for unit testing
])
builder = sdk_container_builder._SdkContainerImageCloudBuilder(options)
builder._temp_src_dir = temp_dir
builder._upload_to_gcs = unittest.mock.MagicMock()

mock_cloudbuild_client = unittest.mock.MagicMock()
mock_build_response = unittest.mock.MagicMock()
mock_build_response.metadata.build.id = 'dummy-id'
mock_build_response.metadata.build.log_url = 'http://dummy.url'
mock_cloudbuild_client.create_build.return_value = mock_build_response
mock_cloudbuild_client.get_build.return_value = cloud_build_types.Build(
status=cloud_build_types.Build.Status.SUCCESS)
builder._cloudbuild_client = mock_cloudbuild_client

builder._invoke_docker_build_and_push(
'gcr.io/test-project/prebuilt_beam_sdk:latest')

mock_cloudbuild_client.create_build.assert_called_once()
create_build_request = mock_cloudbuild_client.create_build.call_args[0][0]
self.assertEqual(
create_build_request.build.options.machine_type,
cloud_build_types.BuildOptions.MachineType.N1_HIGHCPU_8)
Comment on lines +159 to +162

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The assertion relies on create_build being called with positional arguments (call_args[0][0]). If create_build is called using keyword arguments (e.g., create_build(request=...) or create_build(project_id=..., build=...)), this will raise an IndexError.

To make the test more robust against different calling conventions of the Google Cloud client library, consider extracting the build/request object safely from both args and kwargs.

Suggested change
create_build_request = mock_cloudbuild_client.create_build.call_args[0][0]
self.assertEqual(
create_build_request.build.options.machine_type,
cloud_build_types.BuildOptions.MachineType.N1_HIGHCPU_8)
args, kwargs = mock_cloudbuild_client.create_build.call_args
build_obj = (
kwargs.get('build') or
getattr(kwargs.get('request'), 'build', None) or
getattr(args[0], 'build', args[0])
)
self.assertEqual(
build_obj.options.machine_type,
cloud_build_types.BuildOptions.MachineType.N1_HIGHCPU_8)



if __name__ == '__main__':
# Run the tests.
Expand Down
Loading