| # | |
| # 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 pipeline options module.""" | |
| # pytype: skip-file | |
| import json | |
| import logging | |
| import os | |
| import unittest | |
| import hamcrest as hc | |
| import mock | |
| from parameterized import parameterized | |
| from apache_beam.options.pipeline_options import CrossLanguageOptions | |
| from apache_beam.options.pipeline_options import DebugOptions | |
| from apache_beam.options.pipeline_options import GoogleCloudOptions | |
| from apache_beam.options.pipeline_options import JobServerOptions | |
| from apache_beam.options.pipeline_options import PipelineOptions | |
| from apache_beam.options.pipeline_options import ProfilingOptions | |
| from apache_beam.options.pipeline_options import StandardOptions | |
| from apache_beam.options.pipeline_options import TypeOptions | |
| from apache_beam.options.pipeline_options import WorkerOptions | |
| from apache_beam.options.pipeline_options import _BeamArgumentParser | |
| from apache_beam.options.pipeline_options_validator import PipelineOptionsValidator | |
| from apache_beam.options.value_provider import RuntimeValueProvider | |
| from apache_beam.options.value_provider import StaticValueProvider | |
| from apache_beam.transforms.display import DisplayData | |
| from apache_beam.transforms.display_test import DisplayDataItemMatcher | |
| _LOGGER = logging.getLogger(__name__) | |
| try: | |
| from apache_beam.io.gcp.gcsio import GCS_INSTALLED | |
| has_gcsio = GCS_INSTALLED | |
| except ImportError: | |
| has_gcsio = False | |
| # Mock runners to use for validations. | |
| class MockRunners(object): | |
| class DataflowRunner(object): | |
| def get_default_gcp_region(self): | |
| # Return a default so we don't have to specify --region in every test | |
| # (unless specifically testing it). | |
| return 'us-central1' | |
| class MockGoogleCloudOptionsNoBucket(GoogleCloudOptions): | |
| def _create_default_gcs_bucket(self): | |
| return None | |
| class MockGoogleCloudOptionsWithBucket(GoogleCloudOptions): | |
| def _create_default_gcs_bucket(self): | |
| return "gs://default/bucket" | |
| class PipelineOptionsTest(unittest.TestCase): | |
| def setUp(self): | |
| # Reset runtime options to avoid side-effects caused by other tests. | |
| # Note that is_accessible assertions require runtime_options to | |
| # be uninitialized. | |
| RuntimeValueProvider.set_runtime_options(None) | |
| def tearDown(self): | |
| # Reset runtime options to avoid side-effects in other tests. | |
| RuntimeValueProvider.set_runtime_options(None) | |
| TEST_CASES = [ | |
| (['--num_workers', '5'], | |
| { | |
| 'num_workers': 5, | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('num_workers', 5)]), | |
| (['--direct_num_workers', '5'], | |
| { | |
| 'direct_num_workers': 5, | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('direct_num_workers', 5)]), | |
| (['--direct_running_mode', 'multi_threading'], | |
| { | |
| 'direct_running_mode': 'multi_threading', | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('direct_running_mode', 'multi_threading')]), | |
| (['--direct_running_mode', 'multi_processing'], | |
| { | |
| 'direct_running_mode': 'multi_processing', | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('direct_running_mode', 'multi_processing')]), | |
| (['--profile_cpu', '--profile_location', 'gs://bucket/', 'ignored'], | |
| { | |
| 'profile_cpu': True, | |
| 'profile_location': 'gs://bucket/', | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': None | |
| }, | |
| [ | |
| DisplayDataItemMatcher('profile_cpu', True), | |
| DisplayDataItemMatcher('profile_location', 'gs://bucket/') | |
| ]), | |
| (['--num_workers', '5', '--mock_flag'], | |
| { | |
| 'num_workers': 5, | |
| 'mock_flag': True, | |
| 'mock_option': None, | |
| 'mock_multi_option': None | |
| }, | |
| [ | |
| DisplayDataItemMatcher('num_workers', 5), | |
| DisplayDataItemMatcher('mock_flag', True) | |
| ]), | |
| (['--mock_option', 'abc'], { | |
| 'mock_flag': False, 'mock_option': 'abc', 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('mock_option', 'abc')]), | |
| (['--mock_option', ' abc def '], | |
| { | |
| 'mock_flag': False, | |
| 'mock_option': ' abc def ', | |
| 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('mock_option', ' abc def ')]), | |
| (['--mock_option= abc xyz '], | |
| { | |
| 'mock_flag': False, | |
| 'mock_option': ' abc xyz ', | |
| 'mock_multi_option': None | |
| }, [DisplayDataItemMatcher('mock_option', ' abc xyz ')]), | |
| ([ | |
| '--mock_option=gs://my bucket/my folder/my file', | |
| '--mock_multi_option=op1', | |
| '--mock_multi_option=op2' | |
| ], | |
| { | |
| 'mock_flag': False, | |
| 'mock_option': 'gs://my bucket/my folder/my file', | |
| 'mock_multi_option': ['op1', 'op2'] | |
| }, | |
| [ | |
| DisplayDataItemMatcher( | |
| 'mock_option', 'gs://my bucket/my folder/my file'), | |
| DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2']) | |
| ]), | |
| (['--mock_multi_option=op1', '--mock_multi_option=op2'], | |
| { | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': ['op1', 'op2'] | |
| }, [DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2'])]), | |
| (['--mock_json_option={"11a": 0, "37a": 1}'], | |
| { | |
| 'mock_flag': False, | |
| 'mock_option': None, | |
| 'mock_multi_option': None, | |
| 'mock_json_option': { | |
| '11a': 0, '37a': 1 | |
| }, | |
| }, [DisplayDataItemMatcher('mock_json_option', { | |
| '11a': 0, '37a': 1 | |
| })]), | |
| ] | |
| # Used for testing newly added flags. | |
| class MockOptions(PipelineOptions): | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_argument('--mock_flag', action='store_true', help='mock flag') | |
| parser.add_argument('--mock_option', help='mock option') | |
| parser.add_argument( | |
| '--mock_multi_option', action='append', help='mock multi option') | |
| parser.add_argument('--option with space', help='mock option with space') | |
| parser.add_argument('--mock_json_option', type=json.loads, default={}) | |
| # Use with MockOptions in test cases where multiple option classes are needed. | |
| class FakeOptions(PipelineOptions): | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_argument('--fake_flag', action='store_true', help='fake flag') | |
| parser.add_argument('--fake_option', help='fake option') | |
| parser.add_argument( | |
| '--fake_multi_option', action='append', help='fake multi option') | |
| @parameterized.expand(TEST_CASES) | |
| def test_display_data(self, flags, _, display_data): | |
| options = PipelineOptions(flags=flags) | |
| dd = DisplayData.create_from(options) | |
| hc.assert_that(dd.items, hc.contains_inanyorder(*display_data)) | |
| @parameterized.expand(TEST_CASES) | |
| def test_get_all_options_subclass(self, flags, expected, _): | |
| options = PipelineOptionsTest.MockOptions(flags=flags) | |
| self.assertLessEqual(expected.items(), options.get_all_options().items()) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_flag, | |
| expected['mock_flag']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_option, | |
| expected['mock_option']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, | |
| expected['mock_multi_option']) | |
| @parameterized.expand(TEST_CASES) | |
| def test_get_all_options(self, flags, expected, _): | |
| options = PipelineOptions(flags=flags) | |
| self.assertLessEqual(expected.items(), options.get_all_options().items()) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_flag, | |
| expected['mock_flag']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_option, | |
| expected['mock_option']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, | |
| expected['mock_multi_option']) | |
| def test_get_superclass_options(self): | |
| flags = ["--mock_option", "mock", "--fake_option", "fake"] | |
| options = PipelineOptions(flags=flags).view_as( | |
| PipelineOptionsTest.FakeOptions) | |
| items = options.get_all_options(current_only=True).items() | |
| print(items) | |
| self.assertTrue(('fake_option', 'fake') in items) | |
| self.assertFalse(('mock_option', 'mock') in items) | |
| items = options.view_as(PipelineOptionsTest.MockOptions).get_all_options( | |
| current_only=True).items() | |
| self.assertFalse(('fake_option', 'fake') in items) | |
| self.assertTrue(('mock_option', 'mock') in items) | |
| @parameterized.expand(TEST_CASES) | |
| def test_subclasses_of_pipeline_options_can_be_instantiated( | |
| self, flags, expected, _): | |
| mock_options = PipelineOptionsTest.MockOptions(flags=flags) | |
| self.assertEqual(mock_options.mock_flag, expected['mock_flag']) | |
| self.assertEqual(mock_options.mock_option, expected['mock_option']) | |
| self.assertEqual( | |
| mock_options.mock_multi_option, expected['mock_multi_option']) | |
| @parameterized.expand(TEST_CASES) | |
| def test_views_can_be_constructed_from_pipeline_option_subclasses( | |
| self, flags, expected, _): | |
| fake_options = PipelineOptionsTest.FakeOptions(flags=flags) | |
| mock_options = fake_options.view_as(PipelineOptionsTest.MockOptions) | |
| self.assertEqual(mock_options.mock_flag, expected['mock_flag']) | |
| self.assertEqual(mock_options.mock_option, expected['mock_option']) | |
| self.assertEqual( | |
| mock_options.mock_multi_option, expected['mock_multi_option']) | |
| def test_views_do_not_expose_options_defined_by_other_views(self): | |
| flags = ['--mock_option=mock_value', '--fake_option=fake_value'] | |
| options = PipelineOptions(flags) | |
| assert options.view_as( | |
| PipelineOptionsTest.MockOptions).mock_option == 'mock_value' | |
| assert options.view_as( | |
| PipelineOptionsTest.FakeOptions).fake_option == 'fake_value' | |
| assert options.view_as(PipelineOptionsTest.MockOptions).view_as( | |
| PipelineOptionsTest.FakeOptions).fake_option == 'fake_value' | |
| self.assertRaises( | |
| AttributeError, | |
| lambda: options.view_as(PipelineOptionsTest.MockOptions).fake_option) | |
| self.assertRaises( | |
| AttributeError, | |
| lambda: options.view_as(PipelineOptionsTest.MockOptions).view_as( | |
| PipelineOptionsTest.FakeOptions).view_as( | |
| PipelineOptionsTest.MockOptions).fake_option) | |
| @parameterized.expand(TEST_CASES) | |
| def test_from_dictionary(self, flags, expected, _): | |
| options = PipelineOptions(flags=flags) | |
| all_options_dict = options.get_all_options() | |
| options_from_dict = PipelineOptions.from_dictionary(all_options_dict) | |
| self.assertEqual( | |
| options_from_dict.view_as(PipelineOptionsTest.MockOptions).mock_flag, | |
| expected['mock_flag']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_option, | |
| expected['mock_option']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, | |
| expected['mock_multi_option']) | |
| self.assertEqual( | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_json_option, | |
| expected.get('mock_json_option', {})) | |
| def test_none_from_dictionary(self): | |
| class NoneDefaultOptions(PipelineOptions): | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_argument('--test_arg_none', default=None, type=int) | |
| parser.add_argument('--test_arg_int', default=1, type=int) | |
| options_dict = {'test_arg_none': None, 'test_arg_int': 5} | |
| options_from_dict = NoneDefaultOptions.from_dictionary(options_dict) | |
| result = options_from_dict.get_all_options() | |
| self.assertEqual(result['test_arg_int'], 5) | |
| self.assertEqual(result['test_arg_none'], None) | |
| def test_merging_options(self): | |
| opts = PipelineOptions(flags=['--num_workers', '5']) | |
| actual_opts = PipelineOptions.from_runner_api(opts.to_runner_api()) | |
| actual = actual_opts.view_as(WorkerOptions).num_workers | |
| self.assertEqual(5, actual) | |
| def test_merging_options_with_overriden_options(self): | |
| opts = PipelineOptions(flags=['--num_workers', '5']) | |
| base = PipelineOptions(flags=['--num_workers', '2']) | |
| actual_opts = PipelineOptions.from_runner_api(opts.to_runner_api(), base) | |
| actual = actual_opts.view_as(WorkerOptions).num_workers | |
| self.assertEqual(5, actual) | |
| def test_merging_options_with_overriden_runner(self): | |
| opts = PipelineOptions(flags=['--runner', 'FnApiRunner']) | |
| base = PipelineOptions(flags=['--runner', 'Direct']) | |
| actual_opts = PipelineOptions.from_runner_api(opts.to_runner_api(), base) | |
| actual = actual_opts.view_as(StandardOptions).runner | |
| self.assertEqual('Direct', actual) | |
| def test_from_kwargs(self): | |
| class MyOptions(PipelineOptions): | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_argument('--test_arg') | |
| # kwarg takes precedence over parsed flag | |
| options = PipelineOptions(flags=['--test_arg=A'], test_arg='B') | |
| self.assertEqual(options.view_as(MyOptions).test_arg, 'B') | |
| self.assertEqual(options.get_all_options()['test_arg'], 'B') | |
| def test_option_with_space(self): | |
| options = PipelineOptions(flags=['--option with space= value with space']) | |
| self.assertEqual( | |
| getattr( | |
| options.view_as(PipelineOptionsTest.MockOptions), | |
| 'option with space'), | |
| ' value with space') | |
| options_from_dict = PipelineOptions.from_dictionary( | |
| options.get_all_options()) | |
| self.assertEqual( | |
| getattr( | |
| options_from_dict.view_as(PipelineOptionsTest.MockOptions), | |
| 'option with space'), | |
| ' value with space') | |
| def test_retain_unknown_options_binary_store_string(self): | |
| options = PipelineOptions(['--unknown_option', 'some_value']) | |
| result = options.get_all_options(retain_unknown_options=True) | |
| self.assertEqual(result['unknown_option'], 'some_value') | |
| def test_retain_unknown_options_binary_equals_store_string(self): | |
| options = PipelineOptions(['--unknown_option=some_value']) | |
| result = options.get_all_options(retain_unknown_options=True) | |
| self.assertEqual(result['unknown_option'], 'some_value') | |
| def test_retain_unknown_options_binary_multi_equals_store_string(self): | |
| options = PipelineOptions(['--unknown_option=expr = "2 + 2 = 5"']) | |
| result = options.get_all_options(retain_unknown_options=True) | |
| self.assertEqual(result['unknown_option'], 'expr = "2 + 2 = 5"') | |
| def test_retain_unknown_options_binary_single_dash_store_string(self): | |
| options = PipelineOptions(['-i', 'some_value']) | |
| with self.assertRaises(KeyError): | |
| _ = options.get_all_options(retain_unknown_options=True)['i'] | |
| def test_retain_unknown_options_unary_store_true(self): | |
| options = PipelineOptions(['--unknown_option']) | |
| result = options.get_all_options(retain_unknown_options=True) | |
| self.assertEqual(result['unknown_option'], True) | |
| def test_retain_unknown_options_consecutive_unary_store_true(self): | |
| options = PipelineOptions(['--option_foo', '--option_bar']) | |
| result = options.get_all_options(retain_unknown_options=True) | |
| self.assertEqual(result['option_foo'], True) | |
| self.assertEqual(result['option_bar'], True) | |
| def test_retain_unknown_options_unary_single_dash_store_true(self): | |
| options = PipelineOptions(['-i']) | |
| result = options.get_all_options(retain_unknown_options=True) | |
| self.assertEqual(result['i'], True) | |
| def test_override_options(self): | |
| base_flags = ['--num_workers', '5'] | |
| options = PipelineOptions(base_flags) | |
| self.assertEqual(options.get_all_options()['num_workers'], 5) | |
| self.assertEqual(options.get_all_options()['mock_flag'], False) | |
| options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True | |
| self.assertEqual(options.get_all_options()['num_workers'], 5) | |
| self.assertTrue(options.get_all_options()['mock_flag']) | |
| def test_override_init_options(self): | |
| base_flags = ['--num_workers', '5'] | |
| options = PipelineOptions(base_flags, mock_flag=True) | |
| self.assertEqual(options.get_all_options()['num_workers'], 5) | |
| self.assertEqual(options.get_all_options()['mock_flag'], True) | |
| def test_invalid_override_init_options(self): | |
| base_flags = ['--num_workers', '5'] | |
| options = PipelineOptions(base_flags, mock_invalid_flag=True) | |
| self.assertEqual(options.get_all_options()['num_workers'], 5) | |
| self.assertEqual(options.get_all_options()['mock_flag'], False) | |
| def test_experiments(self): | |
| options = PipelineOptions(['--experiment', 'abc', '--experiment', 'def']) | |
| self.assertEqual( | |
| sorted(options.get_all_options()['experiments']), ['abc', 'def']) | |
| options = PipelineOptions(['--experiments', 'abc', '--experiments', 'def']) | |
| self.assertEqual( | |
| sorted(options.get_all_options()['experiments']), ['abc', 'def']) | |
| options = PipelineOptions(flags=['']) | |
| self.assertEqual(options.get_all_options()['experiments'], None) | |
| def test_worker_options(self): | |
| options = PipelineOptions([ | |
| '--machine_type', | |
| 'abc', | |
| '--disk_type', | |
| 'def', | |
| '--disk_provisioned_iops', | |
| '4000', | |
| '--disk_provisioned_throughput_mibps', | |
| '200', | |
| '--element_processing_timeout_minutes', | |
| '10', | |
| ]) | |
| worker_options = options.view_as(WorkerOptions) | |
| self.assertEqual(worker_options.machine_type, 'abc') | |
| self.assertEqual(worker_options.disk_type, 'def') | |
| self.assertEqual(worker_options.disk_provisioned_iops, 4000) | |
| self.assertEqual(worker_options.disk_provisioned_throughput_mibps, 200) | |
| self.assertEqual(worker_options.element_processing_timeout_minutes, 10) | |
| options = PipelineOptions( | |
| ['--worker_machine_type', 'abc', '--worker_disk_type', 'def']) | |
| worker_options = options.view_as(WorkerOptions) | |
| self.assertEqual(worker_options.machine_type, 'abc') | |
| self.assertEqual(worker_options.disk_type, 'def') | |
| def test_beam_services_empty(self): | |
| with mock.patch.dict(os.environ, {}, clear=True): | |
| options = PipelineOptions().view_as(CrossLanguageOptions) | |
| self.assertEqual(options.beam_services, {}) | |
| def test_beam_services_from_env(self): | |
| with mock.patch.dict(os.environ, | |
| {'BEAM_SERVICE_OVERRIDES': '{"foo": "bar"}'}, | |
| clear=True): | |
| options = PipelineOptions().view_as(CrossLanguageOptions) | |
| self.assertEqual(options.beam_services, {'foo': 'bar'}) | |
| def test_beam_services_from_flag(self): | |
| with mock.patch.dict(os.environ, {}, clear=True): | |
| options = PipelineOptions(['--beam_services={"foo": "bar"}' | |
| ]).view_as(CrossLanguageOptions) | |
| self.assertEqual(options.beam_services, {'foo': 'bar'}) | |
| def test_beam_services_from_env_and_flag(self): | |
| with mock.patch.dict( | |
| os.environ, | |
| {'BEAM_SERVICE_OVERRIDES': '{"foo": "bar", "other": "zzz"}'}, | |
| clear=True): | |
| options = PipelineOptions(['--beam_services={"foo": "override"}' | |
| ]).view_as(CrossLanguageOptions) | |
| self.assertEqual( | |
| options.beam_services, { | |
| 'foo': 'override', 'other': 'zzz' | |
| }) | |
| def test_option_modifications_are_shared_between_views(self): | |
| pipeline_options = PipelineOptions([ | |
| '--mock_option', | |
| 'value', | |
| '--mock_flag', | |
| '--mock_multi_option', | |
| 'value1', | |
| '--mock_multi_option', | |
| 'value2', | |
| ]) | |
| mock_options = PipelineOptionsTest.MockOptions([ | |
| '--mock_option', | |
| 'value', | |
| '--mock_flag', | |
| '--mock_multi_option', | |
| 'value1', | |
| '--mock_multi_option', | |
| 'value2', | |
| ]) | |
| for options in [pipeline_options, mock_options]: | |
| view1 = options.view_as(PipelineOptionsTest.MockOptions) | |
| view2 = options.view_as(PipelineOptionsTest.MockOptions) | |
| view1.mock_option = 'new_value' | |
| view1.mock_flag = False | |
| view1.mock_multi_option.append('value3') | |
| view3 = options.view_as(PipelineOptionsTest.MockOptions) | |
| view4 = view1.view_as(PipelineOptionsTest.MockOptions) | |
| view5 = options.view_as(TypeOptions).view_as( | |
| PipelineOptionsTest.MockOptions) | |
| for view in [view1, view2, view3, view4, view5]: | |
| self.assertEqual('new_value', view.mock_option) | |
| self.assertFalse(view.mock_flag) | |
| self.assertEqual(['value1', 'value2', 'value3'], view.mock_multi_option) | |
| def test_uninitialized_option_modifications_are_shared_between_views(self): | |
| options = PipelineOptions([]) | |
| view1 = options.view_as(PipelineOptionsTest.MockOptions) | |
| view2 = options.view_as(PipelineOptionsTest.MockOptions) | |
| view1.mock_option = 'some_value' | |
| view1.mock_flag = False | |
| view1.mock_multi_option = ['value1', 'value2'] | |
| view3 = options.view_as(PipelineOptionsTest.MockOptions) | |
| view4 = view1.view_as(PipelineOptionsTest.MockOptions) | |
| view5 = options.view_as(TypeOptions).view_as( | |
| PipelineOptionsTest.MockOptions) | |
| for view in [view1, view2, view3, view4, view5]: | |
| self.assertEqual('some_value', view.mock_option) | |
| self.assertFalse(view.mock_flag) | |
| self.assertEqual(['value1', 'value2'], view.mock_multi_option) | |
| def test_extra_package(self): | |
| options = PipelineOptions([ | |
| '--extra_package', | |
| 'abc', | |
| '--extra_packages', | |
| 'def', | |
| '--extra_packages', | |
| 'ghi' | |
| ]) | |
| self.assertEqual( | |
| sorted(options.get_all_options()['extra_packages']), | |
| ['abc', 'def', 'ghi']) | |
| options = PipelineOptions(flags=['']) | |
| self.assertEqual(options.get_all_options()['extra_packages'], None) | |
| def test_files_to_stage(self): | |
| options = PipelineOptions([ | |
| '--file_to_stage', | |
| 'abc', | |
| '--files_to_stage', | |
| 'def', | |
| '--files_to_stage', | |
| 'ghi' | |
| ]) | |
| self.assertEqual( | |
| sorted(options.get_all_options()['files_to_stage']), | |
| ['abc', 'def', 'ghi']) | |
| options = PipelineOptions(flags=['']) | |
| self.assertEqual(options.get_all_options()['files_to_stage'], None) | |
| def test_dataflow_job_file(self): | |
| options = PipelineOptions(['--dataflow_job_file', 'abc']) | |
| self.assertEqual(options.get_all_options()['dataflow_job_file'], 'abc') | |
| options = PipelineOptions(flags=['']) | |
| self.assertEqual(options.get_all_options()['dataflow_job_file'], None) | |
| def test_template_location(self): | |
| options = PipelineOptions(['--template_location', 'abc']) | |
| self.assertEqual(options.get_all_options()['template_location'], 'abc') | |
| options = PipelineOptions(flags=['']) | |
| self.assertEqual(options.get_all_options()['template_location'], None) | |
| def test_redefine_options(self): | |
| class TestRedefinedOptions(PipelineOptions): # pylint: disable=unused-variable | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_argument('--redefined_flag', action='store_true') | |
| class TestRedefinedOptions(PipelineOptions): # pylint: disable=function-redefined | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_argument('--redefined_flag', action='store_true') | |
| options = PipelineOptions(['--redefined_flag']) | |
| self.assertTrue(options.get_all_options()['redefined_flag']) | |
| # TODO(https://github.com/apache/beam/issues/18197): Require unique names | |
| # only within a test. For now, <file name acronym>_vp_arg<number> will be | |
| # the convention to name value-provider arguments in tests, as opposed to | |
| # <file name acronym>_non_vp_arg<number> for non-value-provider arguments. | |
| # The number will grow per file as tests are added. | |
| def test_value_provider_options(self): | |
| class UserOptions(PipelineOptions): | |
| @classmethod | |
| def _add_argparse_args(cls, parser): | |
| parser.add_value_provider_argument( | |
| '--pot_vp_arg1', help='This flag is a value provider') | |
| parser.add_value_provider_argument('--pot_vp_arg2', default=1, type=int) | |
| parser.add_argument('--pot_non_vp_arg1', default=1, type=int) | |
| # Provide values: if not provided, the option becomes of the type runtime vp | |
| options = UserOptions(['--pot_vp_arg1', 'hello']) | |
| self.assertIsInstance(options.pot_vp_arg1, StaticValueProvider) | |
| self.assertIsInstance(options.pot_vp_arg2, RuntimeValueProvider) | |
| self.assertIsInstance(options.pot_non_vp_arg1, int) | |
| # Values can be overwritten | |
| options = UserOptions( | |
| pot_vp_arg1=5, | |
| pot_vp_arg2=StaticValueProvider(value_type=str, value='bye'), | |
| pot_non_vp_arg1=RuntimeValueProvider( | |
| option_name='foo', value_type=int, default_value=10)) | |
| self.assertEqual(options.pot_vp_arg1, 5) | |
| self.assertTrue( | |
| options.pot_vp_arg2.is_accessible(), | |
| '%s is not accessible' % options.pot_vp_arg2) | |
| self.assertEqual(options.pot_vp_arg2.get(), 'bye') | |
| self.assertFalse(options.pot_non_vp_arg1.is_accessible()) | |
| with self.assertRaises(RuntimeError): | |
| options.pot_non_vp_arg1.get() | |
| # Converts extra arguments to list value. | |
| def test_extra_args(self): | |
| options = PipelineOptions([ | |
| '--extra_arg', | |
| 'val1', | |
| '--extra_arg', | |
| 'val2', | |
| '--extra_arg=val3', | |
| '--unknown_arg', | |
| 'val4' | |
| ]) | |
| def add_extra_options(parser): | |
| parser.add_argument("--extra_arg", action='append') | |
| self.assertEqual( | |
| options.get_all_options( | |
| add_extra_args_fn=add_extra_options)['extra_arg'], | |
| ['val1', 'val2', 'val3']) | |
| # The argparse package by default tries to autocomplete option names. This | |
| # results in an "ambiguous option" error from argparse when an unknown option | |
| # matching multiple known ones are used. This tests that we suppress this | |
| # error. | |
| def test_unknown_option_prefix(self): | |
| # Test that the "ambiguous option" error is suppressed. | |
| options = PipelineOptions(['--profi', 'val1']) | |
| options.view_as(ProfilingOptions) | |
| # Test that valid errors are not suppressed. | |
| with self.assertRaises(SystemExit): | |
| # Invalid option choice. | |
| options = PipelineOptions(['--type_check_strictness', 'blahblah']) | |
| options.view_as(TypeOptions) | |
| def test_profiling_agent_is_exclusive_with_legacy_profiling_options(self): | |
| options = PipelineOptions(['--profiler_agent=memray']) | |
| validator = PipelineOptionsValidator(options, None) | |
| self.assertEqual(validator.validate(), []) | |
| options = PipelineOptions(['--profiler_agent=memray', '--profile_cpu']) | |
| validator = PipelineOptionsValidator(options, None) | |
| errors = validator.validate() | |
| self.assertTrue( | |
| any('--profiler_agent is mutually exclusive' in err for err in errors)) | |
| options = PipelineOptions(['--profiler_agent=memray', '--profile_memory']) | |
| validator = PipelineOptionsValidator(options, None) | |
| errors = validator.validate() | |
| self.assertTrue( | |
| any('--profiler_agent is mutually exclusive' in err for err in errors)) | |
| def test_profiling_agent_coredump_adds_core_pattern(self): | |
| options = PipelineOptions(['--profiler_agent=coredump']) | |
| validator = PipelineOptionsValidator(options, None) | |
| self.assertEqual(validator.validate(), []) | |
| debug_options = options.view_as(DebugOptions) | |
| self.assertEqual( | |
| debug_options.lookup_experiment('core_pattern'), | |
| '/tmp/beam_coredump.%e.%p') | |
| def test_profile_location_defaulting_and_opt_out(self): | |
| options = PipelineOptions( | |
| ['--profiler_agent=memray', '--temp_location=gs://bucket/temp']) | |
| validator = PipelineOptionsValidator(options, None) | |
| self.assertEqual(validator.validate(), []) | |
| self.assertEqual( | |
| options.view_as(ProfilingOptions).profile_location, | |
| 'gs://bucket/temp/profiles') | |
| options = PipelineOptions([ | |
| '--profiler_agent=memray', | |
| '--temp_location=gs://bucket/temp', | |
| '--profile_location=gs://other-bucket/custom_profiles' | |
| ]) | |
| validator = PipelineOptionsValidator(options, None) | |
| self.assertEqual(validator.validate(), []) | |
| self.assertEqual( | |
| options.view_as(ProfilingOptions).profile_location, | |
| 'gs://other-bucket/custom_profiles') | |
| def test_add_experiment(self): | |
| options = PipelineOptions([]) | |
| options.view_as(DebugOptions).add_experiment('new_experiment') | |
| self.assertEqual(['new_experiment'], | |
| options.view_as(DebugOptions).experiments) | |
| def test_add_experiment_preserves_existing_experiments(self): | |
| options = PipelineOptions(['--experiment=existing_experiment']) | |
| options.view_as(DebugOptions).add_experiment('new_experiment') | |
| self.assertEqual(['existing_experiment', 'new_experiment'], | |
| options.view_as(DebugOptions).experiments) | |
| def test_lookup_experiments(self): | |
| options = PipelineOptions([ | |
| '--experiment=existing_experiment', | |
| '--experiment', | |
| 'key=value', | |
| '--experiment', | |
| 'master_key=k1=v1,k2=v2', | |
| ]) | |
| debug_options = options.view_as(DebugOptions) | |
| self.assertEqual( | |
| 'default_value', | |
| debug_options.lookup_experiment('nonexistent', 'default_value')) | |
| self.assertEqual( | |
| 'value', debug_options.lookup_experiment('key', 'default_value')) | |
| self.assertEqual( | |
| 'k1=v1,k2=v2', debug_options.lookup_experiment('master_key')) | |
| self.assertEqual( | |
| True, debug_options.lookup_experiment('existing_experiment')) | |
| def test_transform_name_mapping(self): | |
| options = PipelineOptions(['--transform_name_mapping={\"from\":\"to\"}']) | |
| mapping = options.view_as(GoogleCloudOptions).transform_name_mapping | |
| self.assertEqual(mapping['from'], 'to') | |
| def test_jar_cache_dir(self): | |
| options = PipelineOptions(['--jar_cache_dir=/path/to/jar_cache_dir']) | |
| jar_cache_dir = options.view_as(JobServerOptions).jar_cache_dir | |
| self.assertEqual(jar_cache_dir, '/path/to/jar_cache_dir') | |
| def test_dataflow_service_options(self): | |
| options = PipelineOptions([ | |
| '--dataflow_service_option', | |
| 'whizz=bang', | |
| '--dataflow_service_option', | |
| 'beep=boop' | |
| ]) | |
| self.assertEqual( | |
| sorted(options.get_all_options()['dataflow_service_options']), | |
| ['beep=boop', 'whizz=bang']) | |
| options = PipelineOptions([ | |
| '--dataflow_service_options', | |
| 'whizz=bang', | |
| '--dataflow_service_options', | |
| 'beep=boop' | |
| ]) | |
| self.assertEqual( | |
| sorted(options.get_all_options()['dataflow_service_options']), | |
| ['beep=boop', 'whizz=bang']) | |
| options = PipelineOptions(flags=['']) | |
| self.assertEqual( | |
| options.get_all_options()['dataflow_service_options'], None) | |
| def test_options_store_false_with_different_dest(self): | |
| parser = _BeamArgumentParser() | |
| for cls in PipelineOptions.__subclasses__(): | |
| cls._add_argparse_args(parser) | |
| actions = parser._actions.copy() | |
| options_to_flags = {} | |
| options_diff_dest_store_true = {} | |
| for i in range(len(actions)): | |
| flag_names = actions[i].option_strings | |
| option_name = actions[i].dest | |
| if isinstance(actions[i].const, bool): | |
| for flag_name in flag_names: | |
| flag_name = flag_name.strip('-') | |
| if flag_name != option_name: | |
| # Capture flags which has store_action=True and has a | |
| # different dest. This behavior would be confusing. | |
| if actions[i].const: | |
| options_diff_dest_store_true[flag_name] = option_name | |
| continue | |
| # check the flags like no_use_public_ips | |
| # default is None, action is {True, False} | |
| if actions[i].default is None: | |
| options_to_flags[option_name] = flag_name | |
| self.assertEqual( | |
| len(options_diff_dest_store_true), | |
| 0, | |
| _LOGGER.error( | |
| "There should be no flags that have a dest " | |
| "different from flag name and action as " | |
| "store_true. It would be confusing " | |
| "to the user. Please specify the dest as the " | |
| "flag_name instead.")) | |
| from apache_beam.options.pipeline_options import _FLAG_THAT_SETS_FALSE_VALUE | |
| self.assertDictEqual( | |
| _FLAG_THAT_SETS_FALSE_VALUE, | |
| options_to_flags, | |
| "If you are adding a new boolean flag with default=None," | |
| " with different dest/option_name from the flag name, please add " | |
| "the dest and the flag name to the map " | |
| "_FLAG_THAT_SETS_FALSE_VALUE in PipelineOptions.py") | |
| def test_gcs_custom_audit_entries(self): | |
| options = PipelineOptions([ | |
| '--gcs_custom_audit_entry=user=test-user', | |
| '--gcs_custom_audit_entry=work=test-work', | |
| '--gcs_custom_audit_entries={"job":"test-job", "id":"1234"}' | |
| ]) | |
| entries = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries | |
| self.assertDictEqual( | |
| entries, | |
| { | |
| 'x-goog-custom-audit-user': 'test-user', | |
| 'x-goog-custom-audit-work': 'test-work', | |
| 'x-goog-custom-audit-job': 'test-job', | |
| 'x-goog-custom-audit-id': '1234' | |
| }) | |
| def test_gcs_custom_audit_entries_wo_duplicated_prefix(self): | |
| options = PipelineOptions([ | |
| '--gcs_custom_audit_entry=x-goog-custom-audit-user=test-user', | |
| '--gcs_custom_audit_entries={"job":"test-job", "id":"1234"}' | |
| ]) | |
| entries = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries | |
| self.assertDictEqual( | |
| entries, | |
| { | |
| 'x-goog-custom-audit-user': 'test-user', | |
| 'x-goog-custom-audit-job': 'test-job', | |
| 'x-goog-custom-audit-id': '1234' | |
| }) | |
| @mock.patch('apache_beam.options.pipeline_options._BeamArgumentParser.error') | |
| def test_gcs_custom_audit_entries_with_errors(self, mock_error): | |
| long_key = 'a' * 65 | |
| options = PipelineOptions([f'--gcs_custom_audit_entry={long_key}=1']) | |
| _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries | |
| self.assertRegex( | |
| mock_error.call_args[0][0], | |
| 'The key .* exceeds the 64-character limit.') | |
| mock_error.reset_mock() | |
| long_value = 'b' * 1201 | |
| options = PipelineOptions([f'--gcs_custom_audit_entry=key={long_value}']) | |
| _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries | |
| self.assertRegex( | |
| mock_error.call_args[0][0], | |
| 'The value .* exceeds the 1200-character limit.') | |
| mock_error.reset_mock() | |
| options = PipelineOptions([ | |
| '--gcs_custom_audit_entry=a=1', | |
| '--gcs_custom_audit_entry=b=2', | |
| '--gcs_custom_audit_entry=c=3', | |
| '--gcs_custom_audit_entry=d=4', | |
| '--gcs_custom_audit_entry=job=test-job' | |
| ]) | |
| _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries | |
| self.assertRegex( | |
| mock_error.call_args[0][0], | |
| 'The maximum allowed number of GCS custom audit entries .*') | |
| mock_error.reset_mock() | |
| options = PipelineOptions([ | |
| '--gcs_custom_audit_entry=a=1', | |
| '--gcs_custom_audit_entry=b=2', | |
| '--gcs_custom_audit_entry=c=3', | |
| '--gcs_custom_audit_entry=d=4' | |
| ]) | |
| _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries | |
| self.assertRegex( | |
| mock_error.call_args[0][0], | |
| 'The maximum allowed number of GCS custom audit entries .*') | |
| def _check_errors(self, options, validator, expected): | |
| if has_gcsio: | |
| with mock.patch('apache_beam.io.gcp.gcsio.GcsIO.is_soft_delete_enabled', | |
| return_value=False): | |
| errors = options._handle_temp_and_staging_locations(validator) | |
| self.assertEqual(errors, expected) | |
| else: | |
| errors = options._handle_temp_and_staging_locations(validator) | |
| self.assertEqual(errors, expected) | |
| def test_validation_good_stg_good_temp(self): | |
| runner = MockRunners.DataflowRunner() | |
| options = GoogleCloudOptions([ | |
| '--project=myproject', | |
| '--staging_location=gs://beam/stg', | |
| '--temp_location=gs://beam/tmp' | |
| ]) | |
| validator = PipelineOptionsValidator(options, runner) | |
| self._check_errors(options, validator, []) | |
| self.assertEqual( | |
| options.get_all_options()['staging_location'], "gs://beam/stg") | |
| self.assertEqual( | |
| options.get_all_options()['temp_location'], "gs://beam/tmp") | |
| def test_validation_bad_stg_good_temp(self): | |
| runner = MockRunners.DataflowRunner() | |
| options = GoogleCloudOptions([ | |
| '--project=myproject', | |
| '--staging_location=badGSpath', | |
| '--temp_location=gs://beam/tmp' | |
| ]) | |
| validator = PipelineOptionsValidator(options, runner) | |
| self._check_errors(options, validator, []) | |
| self.assertEqual( | |
| options.get_all_options()['staging_location'], "gs://beam/tmp") | |
| self.assertEqual( | |
| options.get_all_options()['temp_location'], "gs://beam/tmp") | |
| def test_validation_good_stg_bad_temp(self): | |
| runner = MockRunners.DataflowRunner() | |
| options = GoogleCloudOptions([ | |
| '--project=myproject', | |
| '--staging_location=gs://beam/stg', | |
| '--temp_location=badGSpath' | |
| ]) | |
| validator = PipelineOptionsValidator(options, runner) | |
| self._check_errors(options, validator, []) | |
| self.assertEqual( | |
| options.get_all_options()['staging_location'], "gs://beam/stg") | |
| self.assertEqual( | |
| options.get_all_options()['temp_location'], "gs://beam/stg") | |
| def test_validation_bad_stg_bad_temp_with_default(self): | |
| runner = MockRunners.DataflowRunner() | |
| options = MockGoogleCloudOptionsWithBucket([ | |
| '--project=myproject', | |
| '--staging_location=badGSpath', | |
| '--temp_location=badGSpath' | |
| ]) | |
| validator = PipelineOptionsValidator(options, runner) | |
| self._check_errors(options, validator, []) | |
| self.assertEqual( | |
| options.get_all_options()['staging_location'], "gs://default/bucket") | |
| self.assertEqual( | |
| options.get_all_options()['temp_location'], "gs://default/bucket") | |
| def test_validation_bad_stg_bad_temp_no_default(self): | |
| runner = MockRunners.DataflowRunner() | |
| options = MockGoogleCloudOptionsNoBucket([ | |
| '--project=myproject', | |
| '--staging_location=badGSpath', | |
| '--temp_location=badGSpath' | |
| ]) | |
| validator = PipelineOptionsValidator(options, runner) | |
| self._check_errors( | |
| options, | |
| validator, | |
| [ | |
| 'Invalid GCS path (badGSpath), given for the option: ' \ | |
| 'temp_location.', | |
| 'Invalid GCS path (badGSpath), given for the option: ' \ | |
| 'staging_location.' | |
| ]) | |
| def test_comma_separated_experiments(self): | |
| """Test that comma-separated experiments are parsed correctly.""" | |
| # Test single experiment | |
| options = PipelineOptions(['--experiments=abc']) | |
| self.assertEqual(['abc'], options.get_all_options()['experiments']) | |
| # Test comma-separated experiments | |
| options = PipelineOptions(['--experiments=abc,def,ghi']) | |
| self.assertEqual(['abc', 'def', 'ghi'], | |
| options.get_all_options()['experiments']) | |
| # Test multiple flags with comma-separated values | |
| options = PipelineOptions( | |
| ['--experiments=abc,def', '--experiments=ghi,jkl']) | |
| self.assertEqual(['abc', 'def', 'ghi', 'jkl'], | |
| options.get_all_options()['experiments']) | |
| # Test with spaces around commas | |
| options = PipelineOptions(['--experiments=abc, def , ghi']) | |
| self.assertEqual(['abc', 'def', 'ghi'], | |
| options.get_all_options()['experiments']) | |
| # Test empty values are filtered out | |
| options = PipelineOptions(['--experiments=abc,,def,']) | |
| self.assertEqual(['abc', 'def'], options.get_all_options()['experiments']) | |
| def test_comma_separated_dataflow_service_options(self): | |
| """Test that comma-separated dataflow service options are parsed | |
| correctly.""" | |
| # Test single option | |
| options = PipelineOptions(['--dataflow_service_options=option1=value1']) | |
| self.assertEqual(['option1=value1'], | |
| options.get_all_options()['dataflow_service_options']) | |
| # Test comma-separated options | |
| options = PipelineOptions([ | |
| '--dataflow_service_options=option1=value1,option2=value2,' | |
| 'option3=value3' | |
| ]) | |
| self.assertEqual(['option1=value1', 'option2=value2', 'option3=value3'], | |
| options.get_all_options()['dataflow_service_options']) | |
| # Test multiple flags with comma-separated values | |
| options = PipelineOptions([ | |
| '--dataflow_service_options=option1=value1,option2=value2', | |
| '--dataflow_service_options=option3=value3,option4=value4' | |
| ]) | |
| self.assertEqual([ | |
| 'option1=value1', 'option2=value2', 'option3=value3', 'option4=value4' | |
| ], | |
| options.get_all_options()['dataflow_service_options']) | |
| class CompatVersionTest(unittest.TestCase): | |
| def test_is_compat_version_prior_to(self): | |
| test_cases = [ | |
| # Basic comparison cases | |
| ("1.0.0", "2.0.0", True), # v1 < v2 in major | |
| ("2.0.0", "1.0.0", False), # v1 > v2 in major | |
| ("1.1.0", "1.2.0", True), # v1 < v2 in minor | |
| ("1.2.0", "1.1.0", False), # v1 > v2 in minor | |
| ("1.0.1", "1.0.2", True), # v1 < v2 in patch | |
| ("1.0.2", "1.0.1", False), # v1 > v2 in patch | |
| # Equal versions | |
| ("1.0.0", "1.0.0", False), # Identical | |
| ("0.0.0", "0.0.0", False), # Both zero | |
| # Different lengths - shorter vs longer | |
| ("1.0", "1.0.0", False), # Should be equal (1.0 = 1.0.0) | |
| ("1.0", "1.0.1", True), # 1.0.0 < 1.0.1 | |
| ("1.2", "1.2.0", False), # Should be equal (1.2 = 1.2.0) | |
| ("1.2", "1.2.3", True), # 1.2.0 < 1.2.3 | |
| ("2", "2.0.0", False), # Should be equal (2 = 2.0.0) | |
| ("2", "2.0.1", True), # 2.0.0 < 2.0.1 | |
| ("1", "2.0", True), # 1.0.0 < 2.0.0 | |
| # Different lengths - longer vs shorter | |
| ("1.0.0", "1.0", False), # Should be equal | |
| ("1.0.1", "1.0", False), # 1.0.1 > 1.0.0 | |
| ("1.2.0", "1.2", False), # Should be equal | |
| ("1.2.3", "1.2", False), # 1.2.3 > 1.2.0 | |
| ("2.0.0", "2", False), # Should be equal | |
| ("2.0.1", "2", False), # 2.0.1 > 2.0.0 | |
| ("2.0", "1", False), # 2.0.0 > 1.0.0 | |
| # Mixed length comparisons | |
| ("1.0", "2.0.0", True), # 1.0.0 < 2.0.0 | |
| ("2.0", "1.0.0", False), # 2.0.0 > 1.0.0 | |
| ("1", "1.0.1", True), # 1.0.0 < 1.0.1 | |
| ("1.1", "1.0.9", False), # 1.1.0 > 1.0.9 | |
| # Large numbers | |
| ("1.9.9", "2.0.0", True), # 1.9.9 < 2.0.0 | |
| ("10.0.0", "9.9.9", False), # 10.0.0 > 9.9.9 | |
| ("1.10.0", "1.9.0", False), # 1.10.0 > 1.9.0 | |
| ("1.2.10", "1.2.9", False), # 1.2.10 > 1.2.9 | |
| # Sequential versions | |
| ("1.0.0", "1.0.1", True), | |
| ("1.0.1", "1.0.2", True), | |
| ("1.0.9", "1.1.0", True), | |
| ("1.9.9", "2.0.0", True), | |
| ] | |
| for v1, v2, expected in test_cases: | |
| options = PipelineOptions(update_compatibility_version=v1) | |
| self.assertEqual( | |
| options.is_compat_version_prior_to(v2), | |
| expected, | |
| msg=f"Failed {v1} < {v2} == {expected}") | |
| # None case: no update_compatibility_version set | |
| options_no_compat = PipelineOptions() | |
| self.assertFalse( | |
| options_no_compat.is_compat_version_prior_to("1.0.0"), | |
| msg="Should return False when update_compatibility_version is not set") | |
| if __name__ == '__main__': | |
| logging.getLogger().setLevel(logging.INFO) | |
| unittest.main() |