Skip to content

Commit d60ee55

Browse files
committed
Fix cloud_build_machine_type option
1 parent 0d1f25d commit d60ee55

2 files changed

Lines changed: 37 additions & 1 deletion

File tree

sdks/python/apache_beam/runners/portability/sdk_container_builder.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ def _invoke_docker_build_and_push(self, container_image_name):
243243
build = cloud_build_types.Build()
244244
if self._cloud_build_machine_type:
245245
build.options = cloud_build_types.BuildOptions()
246-
build.options.machineType = self._cloud_build_machine_type
246+
build.options.machine_type = self._cloud_build_machine_type
247247
build.steps = []
248248
step = cloud_build_types.BuildStep()
249249
step.name = 'quay.io/buildah/stable:latest'

sdks/python/apache_beam/runners/portability/sdk_container_builder_test.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import gc
2323
import logging
24+
import tempfile
2425
import unittest
2526
import unittest.mock
2627

@@ -120,6 +121,41 @@ def test_container_image_name_uses_forward_slashes(self):
120121
f' got: {container_image_name}'
121122
)
122123

124+
def test_cloud_builder_assign_machine_type(self):
125+
"""Verify Cloud Build machine_type is correctly set on BuildOptions."""
126+
from google.cloud.devtools.cloudbuild_v1 import types as cloud_build_types
127+
128+
with tempfile.TemporaryDirectory() as temp_dir:
129+
options = pipeline_options.PipelineOptions([
130+
'--project=test-project',
131+
'--temp_location=gs://test-bucket/temp',
132+
'--cloud_build_machine_type=n1-highcpu-8',
133+
'--no_auth', # Avoid fetching credential for unit testing
134+
])
135+
builder = sdk_container_builder._SdkContainerImageCloudBuilder(options)
136+
builder._temp_src_dir = temp_dir
137+
builder._upload_to_gcs = unittest.mock.MagicMock()
138+
139+
mock_cloudbuild_client = unittest.mock.MagicMock()
140+
mock_build_response = unittest.mock.MagicMock()
141+
mock_build_response.metadata.build.id = 'dummy-id'
142+
mock_build_response.metadata.build.log_url = 'http://dummy.url'
143+
mock_cloudbuild_client.create_build.return_value = mock_build_response
144+
mock_cloudbuild_client.get_build.return_value = cloud_build_types.Build(
145+
status=cloud_build_types.Build.Status.SUCCESS)
146+
builder._cloudbuild_client = mock_cloudbuild_client
147+
148+
builder._invoke_docker_build_and_push(
149+
'gcr.io/test-project/prebuilt_beam_sdk:latest')
150+
151+
mock_cloudbuild_client.create_build.assert_called_once()
152+
create_build_request = mock_cloudbuild_client.create_build.call_args[0][0]
153+
self.assertEqual(
154+
create_build_request.build.options.machine_type,
155+
cloud_build_types.BuildOptions.MachineType.N1_HIGHCPU_8)
156+
157+
158+
123159

124160
if __name__ == '__main__':
125161
# Run the tests.

0 commit comments

Comments
 (0)