| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """Unit tests for the apiclient module.""" |
| |
| # pytype: skip-file |
| |
| import itertools |
| import json |
| import logging |
| import os |
| import sys |
| import unittest |
| |
| import mock |
| |
| from apache_beam.io.filesystems import FileSystems |
| from apache_beam.metrics.cells import DistributionData |
| from apache_beam.options.pipeline_options import GoogleCloudOptions |
| from apache_beam.options.pipeline_options import PipelineOptions |
| from apache_beam.pipeline import Pipeline |
| from apache_beam.portability import common_urns |
| from apache_beam.portability.api import beam_runner_api_pb2 |
| from apache_beam.runners.dataflow.internal import names |
| from apache_beam.runners.dataflow.internal.clients import dataflow |
| from apache_beam.transforms import Create |
| from apache_beam.transforms import DataflowDistributionCounter |
| from apache_beam.transforms import DoFn |
| from apache_beam.transforms import ParDo |
| from apache_beam.transforms.environments import DockerEnvironment |
| |
| # Protect against environments where apitools library is not available. |
| # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports |
| try: |
| from apache_beam.runners.dataflow.internal import apiclient |
| except ImportError: |
| apiclient = None # type: ignore |
| # pylint: enable=wrong-import-order, wrong-import-position |
| |
| FAKE_PIPELINE_URL = "gs://invalid-bucket/anywhere" |
| _LOGGER = logging.getLogger(__name__) |
| |
| |
| @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') |
| class UtilTest(unittest.TestCase): |
| @unittest.skip("Enable once BEAM-1080 is fixed.") |
| def test_create_application_client(self): |
| pipeline_options = PipelineOptions() |
| apiclient.DataflowApplicationClient(pipeline_options) |
| |
| def test_pipeline_url(self): |
| pipeline_options = PipelineOptions([ |
| '--subnetwork', |
| '/regions/MY/subnetworks/SUBNETWORK', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| env = apiclient.Environment( |
| [], |
| pipeline_options, |
| '2.0.0', # any environment version |
| FAKE_PIPELINE_URL) |
| |
| recovered_options = None |
| for additionalProperty in env.proto.sdkPipelineOptions.additionalProperties: |
| if additionalProperty.key == 'options': |
| recovered_options = additionalProperty.value |
| break |
| else: |
| self.fail( |
| 'No pipeline options found in %s' % env.proto.sdkPipelineOptions) |
| |
| pipeline_url = None |
| for property in recovered_options.object_value.properties: |
| if property.key == 'pipelineUrl': |
| pipeline_url = property.value |
| break |
| else: |
| self.fail('No pipeline_url found in %s' % recovered_options) |
| |
| self.assertEqual(pipeline_url.string_value, FAKE_PIPELINE_URL) |
| |
| def test_set_network(self): |
| pipeline_options = PipelineOptions([ |
| '--network', |
| 'anetworkname', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual(env.proto.workerPools[0].network, 'anetworkname') |
| |
| def test_set_subnetwork(self): |
| pipeline_options = PipelineOptions([ |
| '--subnetwork', |
| '/regions/MY/subnetworks/SUBNETWORK', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].subnetwork, |
| '/regions/MY/subnetworks/SUBNETWORK') |
| |
| def test_flexrs_blank(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp']) |
| |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual(env.proto.flexResourceSchedulingGoal, None) |
| |
| def test_flexrs_cost(self): |
| pipeline_options = PipelineOptions([ |
| '--flexrs_goal', |
| 'COST_OPTIMIZED', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.flexResourceSchedulingGoal, |
| ( |
| dataflow.Environment.FlexResourceSchedulingGoalValueValuesEnum. |
| FLEXRS_COST_OPTIMIZED)) |
| |
| def test_flexrs_speed(self): |
| pipeline_options = PipelineOptions([ |
| '--flexrs_goal', |
| 'SPEED_OPTIMIZED', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.flexResourceSchedulingGoal, |
| ( |
| dataflow.Environment.FlexResourceSchedulingGoalValueValuesEnum. |
| FLEXRS_SPEED_OPTIMIZED)) |
| |
| def _verify_sdk_harness_container_images_get_set(self, pipeline_options): |
| pipeline = Pipeline(options=pipeline_options) |
| pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned |
| |
| test_environment = DockerEnvironment(container_image='test_default_image') |
| proto_pipeline, _ = pipeline.to_runner_api( |
| return_context=True, default_environment=test_environment) |
| |
| dummy_env = beam_runner_api_pb2.Environment( |
| urn=common_urns.environments.DOCKER.urn, |
| payload=( |
| beam_runner_api_pb2.DockerPayload( |
| container_image='dummy_image')).SerializeToString()) |
| proto_pipeline.components.environments['dummy_env_id'].CopyFrom(dummy_env) |
| |
| dummy_transform = beam_runner_api_pb2.PTransform( |
| environment_id='dummy_env_id') |
| proto_pipeline.components.transforms['dummy_transform_id'].CopyFrom( |
| dummy_transform) |
| |
| env = apiclient.Environment( |
| [], # packages |
| pipeline_options, |
| '2.0.0', # any environment version |
| FAKE_PIPELINE_URL, |
| proto_pipeline) |
| worker_pool = env.proto.workerPools[0] |
| |
| self.assertEqual(2, len(worker_pool.sdkHarnessContainerImages)) |
| |
| env_and_image = [(item.environmentId, item.containerImage) |
| for item in worker_pool.sdkHarnessContainerImages] |
| self.assertIn(('dummy_env_id', 'dummy_image'), env_and_image) |
| self.assertIn((mock.ANY, 'test_default_image'), env_and_image) |
| |
| def test_sdk_harness_container_images_get_set_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--experiments=use_runner_v2', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_sdk_harness_container_images_get_set(pipeline_options) |
| |
| def test_sdk_harness_container_images_get_set_prime(self): |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_options=enable_prime', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_sdk_harness_container_images_get_set(pipeline_options) |
| |
| def _verify_sdk_harness_container_image_overrides(self, pipeline_options): |
| test_environment = DockerEnvironment( |
| container_image='dummy_container_image') |
| proto_pipeline, _ = Pipeline().to_runner_api( |
| return_context=True, default_environment=test_environment) |
| |
| # Accessing non-public method for testing. |
| apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( |
| proto_pipeline, |
| { |
| '.*dummy.*': 'new_dummy_container_image', |
| '.*notfound.*': 'new_dummy_container_image_2' |
| }, |
| pipeline_options) |
| |
| self.assertIsNotNone(1, len(proto_pipeline.components.environments)) |
| env = list(proto_pipeline.components.environments.values())[0] |
| |
| from apache_beam.utils import proto_utils |
| docker_payload = proto_utils.parse_Bytes( |
| env.payload, beam_runner_api_pb2.DockerPayload) |
| |
| # Container image should be overridden by the given override. |
| self.assertEqual( |
| docker_payload.container_image, 'new_dummy_container_image') |
| |
| def test_sdk_harness_container_image_overrides_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--experiments=use_runner_v2', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_sdk_harness_container_image_overrides(pipeline_options) |
| |
| def test_sdk_harness_container_image_overrides_prime(self): |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_options=enable_prime', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_sdk_harness_container_image_overrides(pipeline_options) |
| |
| def _verify_dataflow_container_image_override(self, pipeline_options): |
| pipeline = Pipeline(options=pipeline_options) |
| pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned |
| |
| dummy_env = DockerEnvironment( |
| container_image='apache/beam_dummy_name:dummy_tag') |
| proto_pipeline, _ = pipeline.to_runner_api( |
| return_context=True, default_environment=dummy_env) |
| |
| # Accessing non-public method for testing. |
| apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( |
| proto_pipeline, {}, pipeline_options) |
| |
| from apache_beam.utils import proto_utils |
| found_override = False |
| for env in proto_pipeline.components.environments.values(): |
| docker_payload = proto_utils.parse_Bytes( |
| env.payload, beam_runner_api_pb2.DockerPayload) |
| if docker_payload.container_image.startswith( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): |
| found_override = True |
| |
| self.assertTrue(found_override) |
| |
| def test_dataflow_container_image_override_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--experiments=use_runner_v2', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_dataflow_container_image_override(pipeline_options) |
| |
| def test_dataflow_container_image_override_prime(self): |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_options=enable_prime', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_dataflow_container_image_override(pipeline_options) |
| |
| def _verify_non_apache_container_not_overridden(self, pipeline_options): |
| pipeline = Pipeline(options=pipeline_options) |
| pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned |
| |
| dummy_env = DockerEnvironment( |
| container_image='other_org/dummy_name:dummy_tag') |
| proto_pipeline, _ = pipeline.to_runner_api( |
| return_context=True, default_environment=dummy_env) |
| |
| # Accessing non-public method for testing. |
| apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( |
| proto_pipeline, {}, pipeline_options) |
| |
| self.assertIsNotNone(2, len(proto_pipeline.components.environments)) |
| |
| from apache_beam.utils import proto_utils |
| found_override = False |
| for env in proto_pipeline.components.environments.values(): |
| docker_payload = proto_utils.parse_Bytes( |
| env.payload, beam_runner_api_pb2.DockerPayload) |
| if docker_payload.container_image.startswith( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): |
| found_override = True |
| |
| self.assertFalse(found_override) |
| |
| def test_non_apache_container_not_overridden_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--experiments=use_runner_v2', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_non_apache_container_not_overridden(pipeline_options) |
| |
| def test_non_apache_container_not_overridden_prime(self): |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_options=enable_prime', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| |
| self._verify_non_apache_container_not_overridden(pipeline_options) |
| |
| def _verify_pipeline_sdk_not_overridden(self, pipeline_options): |
| pipeline = Pipeline(options=pipeline_options) |
| pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned |
| |
| proto_pipeline, _ = pipeline.to_runner_api(return_context=True) |
| |
| dummy_env = DockerEnvironment( |
| container_image='dummy_prefix/dummy_name:dummy_tag') |
| proto_pipeline, _ = pipeline.to_runner_api( |
| return_context=True, default_environment=dummy_env) |
| |
| # Accessing non-public method for testing. |
| apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( |
| proto_pipeline, {}, pipeline_options) |
| |
| self.assertIsNotNone(2, len(proto_pipeline.components.environments)) |
| |
| from apache_beam.utils import proto_utils |
| found_override = False |
| for env in proto_pipeline.components.environments.values(): |
| docker_payload = proto_utils.parse_Bytes( |
| env.payload, beam_runner_api_pb2.DockerPayload) |
| if docker_payload.container_image.startswith( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): |
| found_override = True |
| |
| self.assertFalse(found_override) |
| |
| def test_pipeline_sdk_not_overridden_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--experiments=use_runner_v2', |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--sdk_container_image=dummy_prefix/dummy_name:dummy_tag' |
| ]) |
| |
| self._verify_pipeline_sdk_not_overridden(pipeline_options) |
| |
| def test_pipeline_sdk_not_overridden_prime(self): |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_options=enable_prime', |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--sdk_container_image=dummy_prefix/dummy_name:dummy_tag' |
| ]) |
| |
| self._verify_pipeline_sdk_not_overridden(pipeline_options) |
| |
| def test_invalid_default_job_name(self): |
| # Regexp for job names in dataflow. |
| regexp = '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$' |
| |
| job_name = apiclient.Job._build_default_job_name('invalid.-_user_n*/ame') |
| self.assertRegex(job_name, regexp) |
| |
| job_name = apiclient.Job._build_default_job_name( |
| 'invalid-extremely-long.username_that_shouldbeshortened_or_is_invalid') |
| self.assertRegex(job_name, regexp) |
| |
| def test_default_job_name(self): |
| job_name = apiclient.Job.default_job_name(None) |
| regexp = 'beamapp-.*-[0-9]{10}-[0-9]{6}-[a-z0-9]{8}$' |
| self.assertRegex(job_name, regexp) |
| |
| def test_split_int(self): |
| number = 12345 |
| split_number = apiclient.to_split_int(number) |
| self.assertEqual((split_number.lowBits, split_number.highBits), (number, 0)) |
| shift_number = number << 32 |
| split_number = apiclient.to_split_int(shift_number) |
| self.assertEqual((split_number.lowBits, split_number.highBits), (0, number)) |
| |
| def test_translate_distribution_using_accumulator(self): |
| metric_update = dataflow.CounterUpdate() |
| accumulator = mock.Mock() |
| accumulator.min = 1 |
| accumulator.max = 15 |
| accumulator.sum = 16 |
| accumulator.count = 2 |
| apiclient.translate_distribution(accumulator, metric_update) |
| self.assertEqual(metric_update.distribution.min.lowBits, accumulator.min) |
| self.assertEqual(metric_update.distribution.max.lowBits, accumulator.max) |
| self.assertEqual(metric_update.distribution.sum.lowBits, accumulator.sum) |
| self.assertEqual( |
| metric_update.distribution.count.lowBits, accumulator.count) |
| |
| def test_translate_distribution_using_distribution_data(self): |
| metric_update = dataflow.CounterUpdate() |
| distribution_update = DistributionData(16, 2, 1, 15) |
| apiclient.translate_distribution(distribution_update, metric_update) |
| self.assertEqual( |
| metric_update.distribution.min.lowBits, distribution_update.min) |
| self.assertEqual( |
| metric_update.distribution.max.lowBits, distribution_update.max) |
| self.assertEqual( |
| metric_update.distribution.sum.lowBits, distribution_update.sum) |
| self.assertEqual( |
| metric_update.distribution.count.lowBits, distribution_update.count) |
| |
| def test_translate_distribution_using_dataflow_distribution_counter(self): |
| counter_update = DataflowDistributionCounter() |
| counter_update.add_input(1) |
| counter_update.add_input(3) |
| metric_proto = dataflow.CounterUpdate() |
| apiclient.translate_distribution(counter_update, metric_proto) |
| histogram = mock.Mock(firstBucketOffset=None, bucketCounts=None) |
| counter_update.translate_to_histogram(histogram) |
| self.assertEqual(metric_proto.distribution.min.lowBits, counter_update.min) |
| self.assertEqual(metric_proto.distribution.max.lowBits, counter_update.max) |
| self.assertEqual(metric_proto.distribution.sum.lowBits, counter_update.sum) |
| self.assertEqual( |
| metric_proto.distribution.count.lowBits, counter_update.count) |
| self.assertEqual( |
| metric_proto.distribution.histogram.bucketCounts, |
| histogram.bucketCounts) |
| self.assertEqual( |
| metric_proto.distribution.histogram.firstBucketOffset, |
| histogram.firstBucketOffset) |
| |
| def test_translate_means(self): |
| metric_update = dataflow.CounterUpdate() |
| accumulator = mock.Mock() |
| accumulator.sum = 16 |
| accumulator.count = 2 |
| apiclient.MetricUpdateTranslators.translate_scalar_mean_int( |
| accumulator, metric_update) |
| self.assertEqual(metric_update.integerMean.sum.lowBits, accumulator.sum) |
| self.assertEqual(metric_update.integerMean.count.lowBits, accumulator.count) |
| |
| accumulator.sum = 16.0 |
| accumulator.count = 2 |
| apiclient.MetricUpdateTranslators.translate_scalar_mean_float( |
| accumulator, metric_update) |
| self.assertEqual(metric_update.floatingPointMean.sum, accumulator.sum) |
| self.assertEqual( |
| metric_update.floatingPointMean.count.lowBits, accumulator.count) |
| |
| def test_translate_means_using_distribution_accumulator(self): |
| # This is the special case for MeanByteCount. |
| # Which is reported over the FnAPI as a beam distribution, |
| # and to the service as a MetricUpdate IntegerMean. |
| metric_update = dataflow.CounterUpdate() |
| accumulator = mock.Mock() |
| accumulator.min = 7 |
| accumulator.max = 9 |
| accumulator.sum = 16 |
| accumulator.count = 2 |
| apiclient.MetricUpdateTranslators.translate_scalar_mean_int( |
| accumulator, metric_update) |
| self.assertEqual(metric_update.integerMean.sum.lowBits, accumulator.sum) |
| self.assertEqual(metric_update.integerMean.count.lowBits, accumulator.count) |
| |
| accumulator.sum = 16.0 |
| accumulator.count = 2 |
| apiclient.MetricUpdateTranslators.translate_scalar_mean_float( |
| accumulator, metric_update) |
| self.assertEqual(metric_update.floatingPointMean.sum, accumulator.sum) |
| self.assertEqual( |
| metric_update.floatingPointMean.count.lowBits, accumulator.count) |
| |
| def test_default_ip_configuration(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp']) |
| env = apiclient.Environment([], |
| pipeline_options, |
| '2.0.0', |
| FAKE_PIPELINE_URL) |
| self.assertEqual(env.proto.workerPools[0].ipConfiguration, None) |
| |
| def test_public_ip_configuration(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--use_public_ips']) |
| env = apiclient.Environment([], |
| pipeline_options, |
| '2.0.0', |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].ipConfiguration, |
| dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PUBLIC) |
| |
| def test_private_ip_configuration(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--no_use_public_ips']) |
| env = apiclient.Environment([], |
| pipeline_options, |
| '2.0.0', |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].ipConfiguration, |
| dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PRIVATE) |
| |
| def test_number_of_worker_harness_threads(self): |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--number_of_worker_harness_threads', |
| '2' |
| ]) |
| env = apiclient.Environment([], |
| pipeline_options, |
| '2.0.0', |
| FAKE_PIPELINE_URL) |
| self.assertEqual(env.proto.workerPools[0].numThreadsPerWorker, 2) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_harness_override_default_in_released_sdks(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--streaming']) |
| override = ''.join([ |
| 'runner_harness_container_image=', |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY, |
| '/harness:2.2.0' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertIn(override, env.proto.experiments) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_harness_override_absent_in_released_sdks_with_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--streaming', |
| '--experiments=use_runner_v2' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| if env.proto.experiments: |
| for experiment in env.proto.experiments: |
| self.assertNotIn('runner_harness_container_image=', experiment) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_harness_override_custom_in_released_sdks(self): |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--streaming', |
| '--experiments=runner_harness_container_image=fake_image' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| 1, |
| len([ |
| x for x in env.proto.experiments |
| if x.startswith('runner_harness_container_image=') |
| ])) |
| self.assertIn( |
| 'runner_harness_container_image=fake_image', env.proto.experiments) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_harness_override_custom_in_released_sdks_with_runner_v2(self): |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--streaming', |
| '--experiments=runner_harness_container_image=fake_image', |
| '--experiments=use_runner_v2', |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| 1, |
| len([ |
| x for x in env.proto.experiments |
| if x.startswith('runner_harness_container_image=') |
| ])) |
| self.assertIn( |
| 'runner_harness_container_image=fake_image', env.proto.experiments) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0.rc1') |
| def test_harness_override_uses_base_version_in_rc_releases(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--streaming']) |
| override = ''.join([ |
| 'runner_harness_container_image=', |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY, |
| '/harness:2.2.0' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertIn(override, env.proto.experiments) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0.dev') |
| def test_harness_override_absent_in_unreleased_sdk(self): |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--streaming']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| if env.proto.experiments: |
| for experiment in env.proto.experiments: |
| self.assertNotIn('runner_harness_container_image=', experiment) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0.dev') |
| def test_pinned_worker_harness_image_tag_used_in_dev_sdk(self): |
| # streaming, fnapi pipeline. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--streaming']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, |
| ( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/python%d%d-fnapi:%s' % |
| ( |
| sys.version_info[0], |
| sys.version_info[1], |
| names.BEAM_FNAPI_CONTAINER_VERSION))) |
| |
| # batch, legacy pipeline. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, |
| ( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/python%d%d:%s' % ( |
| sys.version_info[0], |
| sys.version_info[1], |
| names.BEAM_CONTAINER_VERSION))) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_worker_harness_image_tag_matches_released_sdk_version(self): |
| # streaming, fnapi pipeline. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--streaming']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, |
| ( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + |
| '/python%d%d-fnapi:2.2.0' % |
| (sys.version_info[0], sys.version_info[1]))) |
| |
| # batch, legacy pipeline. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, |
| ( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/python%d%d:2.2.0' % |
| (sys.version_info[0], sys.version_info[1]))) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0.rc1') |
| def test_worker_harness_image_tag_matches_base_sdk_version_of_an_rc(self): |
| # streaming, fnapi pipeline. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp', '--streaming']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, |
| ( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + |
| '/python%d%d-fnapi:2.2.0' % |
| (sys.version_info[0], sys.version_info[1]))) |
| |
| # batch, legacy pipeline. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, |
| ( |
| names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/python%d%d:2.2.0' % |
| (sys.version_info[0], sys.version_info[1]))) |
| |
| def test_worker_harness_override_takes_precedence_over_sdk_defaults(self): |
| # streaming, fnapi pipeline. |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--streaming', |
| '--sdk_container_image=some:image' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, 'some:image') |
| # batch, legacy pipeline. |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://any-location/temp', |
| '--sdk_container_image=some:image' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.workerPools[0].workerHarnessContainerImage, 'some:image') |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.Job.' |
| 'job_id_for_name', |
| return_value='test_id') |
| def test_transform_name_mapping(self, mock_job): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--update', |
| '--transform_name_mapping', |
| '{\"from\":\"to\"}' |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| self.assertIsNotNone(job.proto.transformNameMapping) |
| |
| def test_created_from_snapshot_id(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--create_from_snapshot', |
| 'test_snapshot_id' |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| self.assertEqual('test_snapshot_id', job.proto.createdFromSnapshotId) |
| |
| def test_labels(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp' |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| self.assertIsNone(job.proto.labels) |
| |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--label', |
| 'key1=value1', |
| '--label', |
| 'key2', |
| '--label', |
| 'key3=value3', |
| '--labels', |
| 'key4=value4', |
| '--labels', |
| 'key5' |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| self.assertEqual(5, len(job.proto.labels.additionalProperties)) |
| self.assertEqual('key1', job.proto.labels.additionalProperties[0].key) |
| self.assertEqual('value1', job.proto.labels.additionalProperties[0].value) |
| self.assertEqual('key2', job.proto.labels.additionalProperties[1].key) |
| self.assertEqual('', job.proto.labels.additionalProperties[1].value) |
| self.assertEqual('key3', job.proto.labels.additionalProperties[2].key) |
| self.assertEqual('value3', job.proto.labels.additionalProperties[2].value) |
| self.assertEqual('key4', job.proto.labels.additionalProperties[3].key) |
| self.assertEqual('value4', job.proto.labels.additionalProperties[3].value) |
| self.assertEqual('key5', job.proto.labels.additionalProperties[4].key) |
| self.assertEqual('', job.proto.labels.additionalProperties[4].value) |
| |
| def test_experiment_use_multiple_sdk_containers(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--experiments', |
| 'beam_fn_api' |
| ]) |
| environment = apiclient.Environment([], |
| pipeline_options, |
| 1, |
| FAKE_PIPELINE_URL) |
| self.assertIn('use_multiple_sdk_containers', environment.proto.experiments) |
| |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--experiments', |
| 'beam_fn_api', |
| '--experiments', |
| 'use_multiple_sdk_containers' |
| ]) |
| environment = apiclient.Environment([], |
| pipeline_options, |
| 1, |
| FAKE_PIPELINE_URL) |
| self.assertIn('use_multiple_sdk_containers', environment.proto.experiments) |
| |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--experiments', |
| 'beam_fn_api', |
| '--experiments', |
| 'no_use_multiple_sdk_containers' |
| ]) |
| environment = apiclient.Environment([], |
| pipeline_options, |
| 1, |
| FAKE_PIPELINE_URL) |
| self.assertNotIn( |
| 'use_multiple_sdk_containers', environment.proto.experiments) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', |
| (3, 8)) |
| def test_get_python_sdk_name(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--experiments', |
| 'beam_fn_api', |
| '--experiments', |
| 'use_multiple_sdk_containers' |
| ]) |
| environment = apiclient.Environment([], |
| pipeline_options, |
| 1, |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| 'Apache Beam Python 3.8 SDK', environment._get_python_sdk_name()) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', |
| (2, 7)) |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_interpreter_version_check_fails_py27(self): |
| pipeline_options = PipelineOptions([]) |
| self.assertRaises( |
| Exception, |
| apiclient._verify_interpreter_version_is_supported, |
| pipeline_options) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', |
| (3, 0, 0)) |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0.dev') |
| def test_interpreter_version_check_passes_on_dev_sdks(self): |
| pipeline_options = PipelineOptions([]) |
| apiclient._verify_interpreter_version_is_supported(pipeline_options) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', |
| (3, 0, 0)) |
| def test_interpreter_version_check_passes_with_experiment(self): |
| pipeline_options = PipelineOptions( |
| ["--experiment=use_unsupported_python_version"]) |
| apiclient._verify_interpreter_version_is_supported(pipeline_options) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', |
| (3, 8, 2)) |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_interpreter_version_check_passes_py38(self): |
| pipeline_options = PipelineOptions([]) |
| apiclient._verify_interpreter_version_is_supported(pipeline_options) |
| |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', |
| (3, 10, 0)) |
| @mock.patch( |
| 'apache_beam.runners.dataflow.internal.apiclient.' |
| 'beam_version.__version__', |
| '2.2.0') |
| def test_interpreter_version_check_fails_on_not_yet_supported_version(self): |
| pipeline_options = PipelineOptions([]) |
| self.assertRaises( |
| Exception, |
| apiclient._verify_interpreter_version_is_supported, |
| pipeline_options) |
| |
| def test_use_unified_worker(self): |
| pipeline_options = PipelineOptions([]) |
| self.assertFalse(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions(['--experiments=beam_fn_api']) |
| self.assertFalse(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions(['--experiments=use_unified_worker']) |
| self.assertTrue(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions( |
| ['--experiments=use_unified_worker', '--experiments=beam_fn_api']) |
| self.assertTrue(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions( |
| ['--experiments=use_runner_v2', '--experiments=beam_fn_api']) |
| self.assertTrue(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions(['--experiments=enable_prime']) |
| self.assertTrue(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions( |
| ['--dataflow_service_options=enable_prime']) |
| self.assertTrue(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_options=enable_prime', |
| '--experiments=disable_prime_runner_v2' |
| ]) |
| self.assertFalse(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions( |
| ['--experiments=enable_prime', '--experiments=disable_prime_runner_v2']) |
| self.assertFalse(apiclient._use_unified_worker(pipeline_options)) |
| |
| pipeline_options = PipelineOptions([ |
| '--experiments=use_unified_worker', |
| '--experiments=use_runner_v2', |
| '--experiments=beam_fn_api' |
| ]) |
| self.assertTrue(apiclient._use_unified_worker(pipeline_options)) |
| |
| def test_get_response_encoding(self): |
| encoding = apiclient.get_response_encoding() |
| |
| assert encoding == 'utf8' |
| |
| def test_graph_is_uploaded(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--experiments', |
| 'beam_fn_api', |
| '--experiments', |
| 'upload_graph' |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| pipeline_options.view_as(GoogleCloudOptions).no_auth = True |
| client = apiclient.DataflowApplicationClient(pipeline_options) |
| with mock.patch.object(client, 'stage_file', side_effect=None): |
| with mock.patch.object(client, 'create_job_description', |
| side_effect=None): |
| with mock.patch.object(client, |
| 'submit_job_description', |
| side_effect=None): |
| client.create_job(job) |
| client.stage_file.assert_called_once_with( |
| mock.ANY, "dataflow_graph.json", mock.ANY) |
| client.create_job_description.assert_called_once() |
| |
| def test_create_job_returns_existing_job(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| self.assertTrue(job.proto.clientRequestId) # asserts non-empty string |
| pipeline_options.view_as(GoogleCloudOptions).no_auth = True |
| client = apiclient.DataflowApplicationClient(pipeline_options) |
| |
| response = dataflow.Job() |
| # different clientRequestId from `job` |
| response.clientRequestId = "20210821081910123456-1234" |
| response.name = 'test_job_name' |
| response.id = '2021-08-19_21_18_43-9756917246311111021' |
| |
| with mock.patch.object(client._client.projects_locations_jobs, |
| 'Create', |
| side_effect=[response]): |
| with mock.patch.object(client, 'create_job_description', |
| side_effect=None): |
| with self.assertRaises( |
| apiclient.DataflowJobAlreadyExistsError) as context: |
| client.create_job(job) |
| |
| self.assertEqual( |
| str(context.exception), |
| 'There is already active job named %s with id: %s. If you want to ' |
| 'submit a second job, try again by setting a different name using ' |
| '--job_name.' % ('test_job_name', response.id)) |
| |
| def test_update_job_returns_existing_job(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--region', |
| 'us-central1', |
| '--update', |
| ]) |
| replace_job_id = '2021-08-21_00_00_01-6081497447916622336' |
| with mock.patch('apache_beam.runners.dataflow.internal.apiclient.Job.' |
| 'job_id_for_name', |
| return_value=replace_job_id) as job_id_for_name_mock: |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| job_id_for_name_mock.assert_called_once() |
| |
| self.assertTrue(job.proto.clientRequestId) # asserts non-empty string |
| |
| pipeline_options.view_as(GoogleCloudOptions).no_auth = True |
| client = apiclient.DataflowApplicationClient(pipeline_options) |
| |
| response = dataflow.Job() |
| # different clientRequestId from `job` |
| response.clientRequestId = "20210821083254123456-1234" |
| response.name = 'test_job_name' |
| response.id = '2021-08-19_21_29_07-5725551945600207770' |
| |
| with mock.patch.object(client, 'create_job_description', side_effect=None): |
| with mock.patch.object(client._client.projects_locations_jobs, |
| 'Create', |
| side_effect=[response]): |
| |
| with self.assertRaises( |
| apiclient.DataflowJobAlreadyExistsError) as context: |
| client.create_job(job) |
| |
| self.assertEqual( |
| str(context.exception), |
| 'The job named %s with id: %s has already been updated into job ' |
| 'id: %s and cannot be updated again.' % |
| ('test_job_name', replace_job_id, response.id)) |
| |
| def test_template_file_generation_with_upload_graph(self): |
| pipeline_options = PipelineOptions([ |
| '--project', |
| 'test_project', |
| '--job_name', |
| 'test_job_name', |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--experiments', |
| 'upload_graph', |
| '--template_location', |
| 'gs://test-location/template' |
| ]) |
| job = apiclient.Job(pipeline_options, FAKE_PIPELINE_URL) |
| job.proto.steps.append(dataflow.Step(name='test_step_name')) |
| |
| pipeline_options.view_as(GoogleCloudOptions).no_auth = True |
| client = apiclient.DataflowApplicationClient(pipeline_options) |
| with mock.patch.object(client, 'stage_file', side_effect=None): |
| with mock.patch.object(client, 'create_job_description', |
| side_effect=None): |
| with mock.patch.object(client, |
| 'submit_job_description', |
| side_effect=None): |
| client.create_job(job) |
| |
| client.stage_file.assert_has_calls([ |
| mock.call(mock.ANY, 'dataflow_graph.json', mock.ANY), |
| mock.call(mock.ANY, 'template', mock.ANY) |
| ]) |
| client.create_job_description.assert_called_once() |
| # template is generated, but job should not be submitted to the |
| # service. |
| client.submit_job_description.assert_not_called() |
| |
| template_filename = client.stage_file.call_args_list[-1][0][1] |
| self.assertTrue('template' in template_filename) |
| template_content = client.stage_file.call_args_list[-1][0][2].read( |
| ).decode('utf-8') |
| template_obj = json.loads(template_content) |
| self.assertFalse(template_obj.get('steps')) |
| self.assertTrue(template_obj['stepsLocation']) |
| |
| def test_stage_resources(self): |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--staging_location', |
| 'gs://test-location/staging', |
| '--no_auth' |
| ]) |
| pipeline = beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments={ |
| 'env1': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/foo1').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/bar1').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/baz').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/renamed1', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed1').SerializeToString()) |
| ]), |
| 'env2': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/foo2').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/bar2').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/baz').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/renamed2', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed2').SerializeToString()) |
| ]) |
| })) |
| client = apiclient.DataflowApplicationClient(pipeline_options) |
| with mock.patch.object(apiclient._LegacyDataflowStager, |
| 'stage_job_resources') as mock_stager: |
| client._stage_resources(pipeline, pipeline_options) |
| mock_stager.assert_called_once_with( |
| [('/tmp/foo1', 'foo1', ''), ('/tmp/bar1', 'bar1', ''), |
| ('/tmp/baz', 'baz1', ''), ('/tmp/renamed1', 'renamed1', 'abcdefg'), |
| ('/tmp/foo2', 'foo2', ''), ('/tmp/bar2', 'bar2', '')], |
| staging_location='gs://test-location/staging') |
| |
| pipeline_expected = beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments={ |
| 'env1': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/foo1' |
| ).SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/bar1'). |
| SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/baz1'). |
| SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/renamed1', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed1').SerializeToString()) |
| ]), |
| 'env2': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/foo2'). |
| SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/bar2'). |
| SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/baz1'). |
| SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/renamed1', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed1').SerializeToString()) |
| ]) |
| })) |
| self.assertEqual(pipeline, pipeline_expected) |
| |
| def test_set_dataflow_service_option(self): |
| pipeline_options = PipelineOptions([ |
| '--dataflow_service_option', |
| 'whizz=bang', |
| '--temp_location', |
| 'gs://any-location/temp' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual(env.proto.serviceOptions, ['whizz=bang']) |
| |
| def test_enable_hot_key_logging(self): |
| # Tests that the enable_hot_key_logging is not set by default. |
| pipeline_options = PipelineOptions( |
| ['--temp_location', 'gs://any-location/temp']) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertIsNone(env.proto.debugOptions) |
| |
| # Now test that it is set when given. |
| pipeline_options = PipelineOptions([ |
| '--enable_hot_key_logging', '--temp_location', 'gs://any-location/temp' |
| ]) |
| env = apiclient.Environment( |
| [], #packages |
| pipeline_options, |
| '2.0.0', #any environment version |
| FAKE_PIPELINE_URL) |
| self.assertEqual( |
| env.proto.debugOptions, dataflow.DebugOptions(enableHotKeyLogging=True)) |
| |
| def _mock_uncached_copy(self, staging_root, src, sha256, dst_name=None): |
| sha_prefix = sha256[0:2] |
| gcs_cache_path = FileSystems.join( |
| staging_root, |
| apiclient.DataflowApplicationClient._GCS_CACHE_PREFIX, |
| sha_prefix, |
| sha256) |
| |
| if not dst_name: |
| _, dst_name = os.path.split(src) |
| return [ |
| mock.call.gcs_exists(gcs_cache_path), |
| mock.call.gcs_upload(src, gcs_cache_path), |
| mock.call.gcs_gcs_copy( |
| source_file_names=[gcs_cache_path], |
| destination_file_names=[f'gs://test-location/staging/{dst_name}']) |
| ] |
| |
| def _mock_cached_copy(self, staging_root, src, sha256, dst_name=None): |
| uncached = self._mock_uncached_copy(staging_root, src, sha256, dst_name) |
| uncached.pop(1) |
| return uncached |
| |
| def test_stage_artifacts_with_caching(self): |
| pipeline_options = PipelineOptions([ |
| '--temp_location', |
| 'gs://test-location/temp', |
| '--staging_location', |
| 'gs://test-location/staging', |
| '--no_auth', |
| '--enable_artifact_caching' |
| ]) |
| pipeline = beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments={ |
| 'env1': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/foo1', |
| sha256='abcd').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/bar1', |
| sha256='defg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload(path='/tmp/baz', sha256='hijk' |
| ).SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/renamed1', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed1').SerializeToString()) |
| ]), |
| 'env2': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/foo2', |
| sha256='lmno').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/bar2', |
| sha256='pqrs').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload(path='/tmp/baz', sha256='tuv' |
| ).SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2. |
| ArtifactFilePayload( |
| path='/tmp/renamed2', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed2').SerializeToString()) |
| ]) |
| })) |
| client = apiclient.DataflowApplicationClient(pipeline_options) |
| staging_root = 'gs://test-location/staging' |
| |
| # every other artifact already exists |
| n = [0] |
| |
| def exists_return_value(*args): |
| n[0] += 1 |
| return n[0] % 2 == 0 |
| |
| with mock.patch.object(FileSystems, |
| 'exists', |
| side_effect=exists_return_value) as mock_gcs_exists: |
| with mock.patch.object(apiclient.DataflowApplicationClient, |
| '_uncached_gcs_file_copy') as mock_gcs_copy: |
| with mock.patch.object(FileSystems, 'copy') as mock_gcs_gcs_copy: |
| |
| manager = mock.Mock() |
| manager.attach_mock(mock_gcs_exists, 'gcs_exists') |
| manager.attach_mock(mock_gcs_copy, 'gcs_upload') |
| manager.attach_mock(mock_gcs_gcs_copy, 'gcs_gcs_copy') |
| |
| client._stage_resources(pipeline, pipeline_options) |
| expected_calls = list( |
| itertools.chain.from_iterable([ |
| self._mock_uncached_copy(staging_root, '/tmp/foo1', 'abcd'), |
| self._mock_cached_copy(staging_root, '/tmp/bar1', 'defg'), |
| self._mock_uncached_copy( |
| staging_root, '/tmp/baz', 'hijk', 'baz1'), |
| self._mock_cached_copy( |
| staging_root, '/tmp/renamed1', 'abcdefg'), |
| self._mock_uncached_copy(staging_root, '/tmp/foo2', 'lmno'), |
| self._mock_cached_copy(staging_root, '/tmp/bar2', 'pqrs'), |
| ])) |
| assert manager.mock_calls == expected_calls |
| |
| pipeline_expected = beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments={ |
| 'env1': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/foo1', |
| sha256='abcd').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/bar1', |
| sha256='defg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/baz1', |
| sha256='hijk').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/renamed1', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed1').SerializeToString()) |
| ]), |
| 'env2': beam_runner_api_pb2.Environment( |
| dependencies=[ |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/foo2', |
| sha256='lmno').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='foo2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/bar2', |
| sha256='pqrs').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='bar2').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/baz1', |
| sha256='tuv').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='baz1').SerializeToString()), |
| beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.URL.urn, |
| type_payload=beam_runner_api_pb2.ArtifactUrlPayload( |
| url='gs://test-location/staging/renamed1', |
| sha256='abcdefg').SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2. |
| ArtifactStagingToRolePayload( |
| staged_name='renamed1').SerializeToString()) |
| ]) |
| })) |
| self.assertEqual(pipeline, pipeline_expected) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
