gemini-code-assist[bot] commented on code in PR #39342:
URL: https://github.com/apache/beam/pull/39342#discussion_r3588495227
##########
sdks/python/apache_beam/runners/portability/sdk_container_builder_test.py:
##########
@@ -120,6 +121,41 @@ def test_container_image_name_uses_forward_slashes(self):
f' got: {container_image_name}'
)
+ 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)
+
+
+
Review Comment:
There are too many blank lines added at the end of the test method.
According to PEP 8, there should be exactly two blank lines between top-level
definitions (such as the class containing this test and the `if __name__ ==
'__main__':` block). Please remove the extra blank lines.
```suggestion
```
<details>
<summary>References</summary>
1. PEP 8: Blank Lines. Two blank lines are expected before top-level
definitions.
<sup>([link](https://github.com/apache/beam/blob/master/.gemini/styleguide.md))</sup>
</details>
##########
sdks/python/apache_beam/runners/portability/sdk_container_builder_test.py:
##########
@@ -120,6 +121,41 @@ def test_container_image_name_uses_forward_slashes(self):
f' got: {container_image_name}'
)
+ 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)
Review Comment:
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`.
```suggestion
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)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
