github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/options/pipeline_options_test.py (about) 1 # 2 # Licensed to the Apache Software Foundation (ASF) under one or more 3 # contributor license agreements. See the NOTICE file distributed with 4 # this work for additional information regarding copyright ownership. 5 # The ASF licenses this file to You under the Apache License, Version 2.0 6 # (the "License"); you may not use this file except in compliance with 7 # the License. You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 # 17 18 """Unit tests for the pipeline options module.""" 19 20 # pytype: skip-file 21 22 import json 23 import logging 24 import unittest 25 26 import hamcrest as hc 27 28 from apache_beam.options.pipeline_options import DebugOptions 29 from apache_beam.options.pipeline_options import GoogleCloudOptions 30 from apache_beam.options.pipeline_options import PipelineOptions 31 from apache_beam.options.pipeline_options import ProfilingOptions 32 from apache_beam.options.pipeline_options import TypeOptions 33 from apache_beam.options.pipeline_options import WorkerOptions 34 from apache_beam.options.pipeline_options import _BeamArgumentParser 35 from apache_beam.options.value_provider import RuntimeValueProvider 36 from apache_beam.options.value_provider import StaticValueProvider 37 from apache_beam.transforms.display import DisplayData 38 from apache_beam.transforms.display_test import DisplayDataItemMatcher 39 40 _LOGGER = logging.getLogger(__name__) 41 42 43 class PipelineOptionsTest(unittest.TestCase): 44 def setUp(self): 45 # Reset runtime options to avoid side-effects caused by other tests. 46 # Note that is_accessible assertions require runtime_options to 47 # be uninitialized. 48 RuntimeValueProvider.set_runtime_options(None) 49 50 def tearDown(self): 51 # Reset runtime options to avoid side-effects in other tests. 52 RuntimeValueProvider.set_runtime_options(None) 53 54 TEST_CASES = [ 55 { 56 'flags': ['--num_workers', '5'], 57 'expected': { 58 'num_workers': 5, 59 'mock_flag': False, 60 'mock_option': None, 61 'mock_multi_option': None 62 }, 63 'display_data': [DisplayDataItemMatcher('num_workers', 5)] 64 }, 65 { 66 'flags': ['--direct_num_workers', '5'], 67 'expected': { 68 'direct_num_workers': 5, 69 'mock_flag': False, 70 'mock_option': None, 71 'mock_multi_option': None 72 }, 73 'display_data': [DisplayDataItemMatcher('direct_num_workers', 5)] 74 }, 75 { 76 'flags': ['--direct_running_mode', 'multi_threading'], 77 'expected': { 78 'direct_running_mode': 'multi_threading', 79 'mock_flag': False, 80 'mock_option': None, 81 'mock_multi_option': None 82 }, 83 'display_data': [ 84 DisplayDataItemMatcher('direct_running_mode', 'multi_threading') 85 ] 86 }, 87 { 88 'flags': ['--direct_running_mode', 'multi_processing'], 89 'expected': { 90 'direct_running_mode': 'multi_processing', 91 'mock_flag': False, 92 'mock_option': None, 93 'mock_multi_option': None 94 }, 95 'display_data': [ 96 DisplayDataItemMatcher('direct_running_mode', 'multi_processing') 97 ] 98 }, 99 { 100 'flags': [ 101 '--profile_cpu', '--profile_location', 'gs://bucket/', 'ignored' 102 ], 103 'expected': { 104 'profile_cpu': True, 105 'profile_location': 'gs://bucket/', 106 'mock_flag': False, 107 'mock_option': None, 108 'mock_multi_option': None 109 }, 110 'display_data': [ 111 DisplayDataItemMatcher('profile_cpu', True), 112 DisplayDataItemMatcher('profile_location', 'gs://bucket/') 113 ] 114 }, 115 { 116 'flags': ['--num_workers', '5', '--mock_flag'], 117 'expected': { 118 'num_workers': 5, 119 'mock_flag': True, 120 'mock_option': None, 121 'mock_multi_option': None 122 }, 123 'display_data': [ 124 DisplayDataItemMatcher('num_workers', 5), 125 DisplayDataItemMatcher('mock_flag', True) 126 ] 127 }, 128 { 129 'flags': ['--mock_option', 'abc'], 130 'expected': { 131 'mock_flag': False, 132 'mock_option': 'abc', 133 'mock_multi_option': None 134 }, 135 'display_data': [DisplayDataItemMatcher('mock_option', 'abc')] 136 }, 137 { 138 'flags': ['--mock_option', ' abc def '], 139 'expected': { 140 'mock_flag': False, 141 'mock_option': ' abc def ', 142 'mock_multi_option': None 143 }, 144 'display_data': [DisplayDataItemMatcher('mock_option', ' abc def ')] 145 }, 146 { 147 'flags': ['--mock_option= abc xyz '], 148 'expected': { 149 'mock_flag': False, 150 'mock_option': ' abc xyz ', 151 'mock_multi_option': None 152 }, 153 'display_data': [DisplayDataItemMatcher('mock_option', ' abc xyz ')] 154 }, 155 { 156 'flags': [ 157 '--mock_option=gs://my bucket/my folder/my file', 158 '--mock_multi_option=op1', 159 '--mock_multi_option=op2' 160 ], 161 'expected': { 162 'mock_flag': False, 163 'mock_option': 'gs://my bucket/my folder/my file', 164 'mock_multi_option': ['op1', 'op2'] 165 }, 166 'display_data': [ 167 DisplayDataItemMatcher( 168 'mock_option', 'gs://my bucket/my folder/my file'), 169 DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2']) 170 ] 171 }, 172 { 173 'flags': ['--mock_multi_option=op1', '--mock_multi_option=op2'], 174 'expected': { 175 'mock_flag': False, 176 'mock_option': None, 177 'mock_multi_option': ['op1', 'op2'] 178 }, 179 'display_data': [ 180 DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2']) 181 ] 182 }, 183 { 184 'flags': ['--mock_json_option={"11a": 0, "37a": 1}'], 185 'expected': { 186 'mock_flag': False, 187 'mock_option': None, 188 'mock_multi_option': None, 189 'mock_json_option': { 190 '11a': 0, '37a': 1 191 }, 192 }, 193 'display_data': [ 194 DisplayDataItemMatcher('mock_json_option', { 195 '11a': 0, '37a': 1 196 }) 197 ] 198 }, 199 ] 200 201 # Used for testing newly added flags. 202 class MockOptions(PipelineOptions): 203 @classmethod 204 def _add_argparse_args(cls, parser): 205 parser.add_argument('--mock_flag', action='store_true', help='mock flag') 206 parser.add_argument('--mock_option', help='mock option') 207 parser.add_argument( 208 '--mock_multi_option', action='append', help='mock multi option') 209 parser.add_argument('--option with space', help='mock option with space') 210 parser.add_argument('--mock_json_option', type=json.loads, default={}) 211 212 # Use with MockOptions in test cases where multiple option classes are needed. 213 class FakeOptions(PipelineOptions): 214 @classmethod 215 def _add_argparse_args(cls, parser): 216 parser.add_argument('--fake_flag', action='store_true', help='fake flag') 217 parser.add_argument('--fake_option', help='fake option') 218 parser.add_argument( 219 '--fake_multi_option', action='append', help='fake multi option') 220 221 def test_display_data(self): 222 for case in PipelineOptionsTest.TEST_CASES: 223 options = PipelineOptions(flags=case['flags']) 224 dd = DisplayData.create_from(options) 225 hc.assert_that(dd.items, hc.contains_inanyorder(*case['display_data'])) 226 227 def test_get_all_options_subclass(self): 228 for case in PipelineOptionsTest.TEST_CASES: 229 options = PipelineOptionsTest.MockOptions(flags=case['flags']) 230 self.assertDictContainsSubset(case['expected'], options.get_all_options()) 231 self.assertEqual( 232 options.view_as(PipelineOptionsTest.MockOptions).mock_flag, 233 case['expected']['mock_flag']) 234 self.assertEqual( 235 options.view_as(PipelineOptionsTest.MockOptions).mock_option, 236 case['expected']['mock_option']) 237 self.assertEqual( 238 options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, 239 case['expected']['mock_multi_option']) 240 241 def test_get_all_options(self): 242 for case in PipelineOptionsTest.TEST_CASES: 243 options = PipelineOptions(flags=case['flags']) 244 self.assertDictContainsSubset(case['expected'], options.get_all_options()) 245 self.assertEqual( 246 options.view_as(PipelineOptionsTest.MockOptions).mock_flag, 247 case['expected']['mock_flag']) 248 self.assertEqual( 249 options.view_as(PipelineOptionsTest.MockOptions).mock_option, 250 case['expected']['mock_option']) 251 self.assertEqual( 252 options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, 253 case['expected']['mock_multi_option']) 254 255 def test_sublcalsses_of_pipeline_options_can_be_instantiated(self): 256 for case in PipelineOptionsTest.TEST_CASES: 257 mock_options = PipelineOptionsTest.MockOptions(flags=case['flags']) 258 self.assertEqual(mock_options.mock_flag, case['expected']['mock_flag']) 259 self.assertEqual( 260 mock_options.mock_option, case['expected']['mock_option']) 261 self.assertEqual( 262 mock_options.mock_multi_option, case['expected']['mock_multi_option']) 263 264 def test_views_can_be_constructed_from_pipeline_option_subclasses(self): 265 for case in PipelineOptionsTest.TEST_CASES: 266 fake_options = PipelineOptionsTest.FakeOptions(flags=case['flags']) 267 mock_options = fake_options.view_as(PipelineOptionsTest.MockOptions) 268 269 self.assertEqual(mock_options.mock_flag, case['expected']['mock_flag']) 270 self.assertEqual( 271 mock_options.mock_option, case['expected']['mock_option']) 272 self.assertEqual( 273 mock_options.mock_multi_option, case['expected']['mock_multi_option']) 274 275 def test_views_do_not_expose_options_defined_by_other_views(self): 276 flags = ['--mock_option=mock_value', '--fake_option=fake_value'] 277 278 options = PipelineOptions(flags) 279 assert options.view_as( 280 PipelineOptionsTest.MockOptions).mock_option == 'mock_value' 281 assert options.view_as( 282 PipelineOptionsTest.FakeOptions).fake_option == 'fake_value' 283 assert options.view_as(PipelineOptionsTest.MockOptions).view_as( 284 PipelineOptionsTest.FakeOptions).fake_option == 'fake_value' 285 286 self.assertRaises( 287 AttributeError, 288 lambda: options.view_as(PipelineOptionsTest.MockOptions).fake_option) 289 self.assertRaises( 290 AttributeError, 291 lambda: options.view_as(PipelineOptionsTest.MockOptions).view_as( 292 PipelineOptionsTest.FakeOptions).view_as( 293 PipelineOptionsTest.MockOptions).fake_option) 294 295 def test_from_dictionary(self): 296 for case in PipelineOptionsTest.TEST_CASES: 297 options = PipelineOptions(flags=case['flags']) 298 all_options_dict = options.get_all_options() 299 options_from_dict = PipelineOptions.from_dictionary(all_options_dict) 300 self.assertEqual( 301 options_from_dict.view_as(PipelineOptionsTest.MockOptions).mock_flag, 302 case['expected']['mock_flag']) 303 self.assertEqual( 304 options.view_as(PipelineOptionsTest.MockOptions).mock_option, 305 case['expected']['mock_option']) 306 self.assertEqual( 307 options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, 308 case['expected']['mock_multi_option']) 309 self.assertEqual( 310 options.view_as(PipelineOptionsTest.MockOptions).mock_json_option, 311 case['expected'].get('mock_json_option', {})) 312 313 def test_none_from_dictionary(self): 314 class NoneDefaultOptions(PipelineOptions): 315 @classmethod 316 def _add_argparse_args(cls, parser): 317 parser.add_argument('--test_arg_none', default=None, type=int) 318 parser.add_argument('--test_arg_int', default=1, type=int) 319 320 options_dict = {'test_arg_none': None, 'test_arg_int': 5} 321 options_from_dict = NoneDefaultOptions.from_dictionary(options_dict) 322 result = options_from_dict.get_all_options() 323 self.assertEqual(result['test_arg_int'], 5) 324 self.assertEqual(result['test_arg_none'], None) 325 326 def test_option_with_space(self): 327 options = PipelineOptions(flags=['--option with space= value with space']) 328 self.assertEqual( 329 getattr( 330 options.view_as(PipelineOptionsTest.MockOptions), 331 'option with space'), 332 ' value with space') 333 options_from_dict = PipelineOptions.from_dictionary( 334 options.get_all_options()) 335 self.assertEqual( 336 getattr( 337 options_from_dict.view_as(PipelineOptionsTest.MockOptions), 338 'option with space'), 339 ' value with space') 340 341 def test_retain_unknown_options_binary_store_string(self): 342 options = PipelineOptions(['--unknown_option', 'some_value']) 343 result = options.get_all_options(retain_unknown_options=True) 344 self.assertEqual(result['unknown_option'], 'some_value') 345 346 def test_retain_unknown_options_binary_equals_store_string(self): 347 options = PipelineOptions(['--unknown_option=some_value']) 348 result = options.get_all_options(retain_unknown_options=True) 349 self.assertEqual(result['unknown_option'], 'some_value') 350 351 def test_retain_unknown_options_binary_multi_equals_store_string(self): 352 options = PipelineOptions(['--unknown_option=expr = "2 + 2 = 5"']) 353 result = options.get_all_options(retain_unknown_options=True) 354 self.assertEqual(result['unknown_option'], 'expr = "2 + 2 = 5"') 355 356 def test_retain_unknown_options_binary_single_dash_store_string(self): 357 options = PipelineOptions(['-i', 'some_value']) 358 with self.assertRaises(KeyError): 359 _ = options.get_all_options(retain_unknown_options=True)['i'] 360 361 def test_retain_unknown_options_unary_store_true(self): 362 options = PipelineOptions(['--unknown_option']) 363 result = options.get_all_options(retain_unknown_options=True) 364 self.assertEqual(result['unknown_option'], True) 365 366 def test_retain_unknown_options_consecutive_unary_store_true(self): 367 options = PipelineOptions(['--option_foo', '--option_bar']) 368 result = options.get_all_options(retain_unknown_options=True) 369 self.assertEqual(result['option_foo'], True) 370 self.assertEqual(result['option_bar'], True) 371 372 def test_retain_unknown_options_unary_single_dash_store_true(self): 373 options = PipelineOptions(['-i']) 374 result = options.get_all_options(retain_unknown_options=True) 375 self.assertEqual(result['i'], True) 376 377 def test_override_options(self): 378 base_flags = ['--num_workers', '5'] 379 options = PipelineOptions(base_flags) 380 self.assertEqual(options.get_all_options()['num_workers'], 5) 381 self.assertEqual(options.get_all_options()['mock_flag'], False) 382 383 options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True 384 self.assertEqual(options.get_all_options()['num_workers'], 5) 385 self.assertTrue(options.get_all_options()['mock_flag']) 386 387 def test_override_init_options(self): 388 base_flags = ['--num_workers', '5'] 389 options = PipelineOptions(base_flags, mock_flag=True) 390 self.assertEqual(options.get_all_options()['num_workers'], 5) 391 self.assertEqual(options.get_all_options()['mock_flag'], True) 392 393 def test_invalid_override_init_options(self): 394 base_flags = ['--num_workers', '5'] 395 options = PipelineOptions(base_flags, mock_invalid_flag=True) 396 self.assertEqual(options.get_all_options()['num_workers'], 5) 397 self.assertEqual(options.get_all_options()['mock_flag'], False) 398 399 def test_experiments(self): 400 options = PipelineOptions(['--experiment', 'abc', '--experiment', 'def']) 401 self.assertEqual( 402 sorted(options.get_all_options()['experiments']), ['abc', 'def']) 403 404 options = PipelineOptions(['--experiments', 'abc', '--experiments', 'def']) 405 self.assertEqual( 406 sorted(options.get_all_options()['experiments']), ['abc', 'def']) 407 408 options = PipelineOptions(flags=['']) 409 self.assertEqual(options.get_all_options()['experiments'], None) 410 411 def test_worker_options(self): 412 options = PipelineOptions(['--machine_type', 'abc', '--disk_type', 'def']) 413 worker_options = options.view_as(WorkerOptions) 414 self.assertEqual(worker_options.machine_type, 'abc') 415 self.assertEqual(worker_options.disk_type, 'def') 416 417 options = PipelineOptions( 418 ['--worker_machine_type', 'abc', '--worker_disk_type', 'def']) 419 worker_options = options.view_as(WorkerOptions) 420 self.assertEqual(worker_options.machine_type, 'abc') 421 self.assertEqual(worker_options.disk_type, 'def') 422 423 def test_option_modifications_are_shared_between_views(self): 424 pipeline_options = PipelineOptions([ 425 '--mock_option', 426 'value', 427 '--mock_flag', 428 '--mock_multi_option', 429 'value1', 430 '--mock_multi_option', 431 'value2', 432 ]) 433 434 mock_options = PipelineOptionsTest.MockOptions([ 435 '--mock_option', 436 'value', 437 '--mock_flag', 438 '--mock_multi_option', 439 'value1', 440 '--mock_multi_option', 441 'value2', 442 ]) 443 444 for options in [pipeline_options, mock_options]: 445 view1 = options.view_as(PipelineOptionsTest.MockOptions) 446 view2 = options.view_as(PipelineOptionsTest.MockOptions) 447 448 view1.mock_option = 'new_value' 449 view1.mock_flag = False 450 view1.mock_multi_option.append('value3') 451 452 view3 = options.view_as(PipelineOptionsTest.MockOptions) 453 view4 = view1.view_as(PipelineOptionsTest.MockOptions) 454 view5 = options.view_as(TypeOptions).view_as( 455 PipelineOptionsTest.MockOptions) 456 457 for view in [view1, view2, view3, view4, view5]: 458 self.assertEqual('new_value', view.mock_option) 459 self.assertFalse(view.mock_flag) 460 self.assertEqual(['value1', 'value2', 'value3'], view.mock_multi_option) 461 462 def test_uninitialized_option_modifications_are_shared_between_views(self): 463 options = PipelineOptions([]) 464 465 view1 = options.view_as(PipelineOptionsTest.MockOptions) 466 view2 = options.view_as(PipelineOptionsTest.MockOptions) 467 468 view1.mock_option = 'some_value' 469 view1.mock_flag = False 470 view1.mock_multi_option = ['value1', 'value2'] 471 472 view3 = options.view_as(PipelineOptionsTest.MockOptions) 473 view4 = view1.view_as(PipelineOptionsTest.MockOptions) 474 view5 = options.view_as(TypeOptions).view_as( 475 PipelineOptionsTest.MockOptions) 476 477 for view in [view1, view2, view3, view4, view5]: 478 self.assertEqual('some_value', view.mock_option) 479 self.assertFalse(view.mock_flag) 480 self.assertEqual(['value1', 'value2'], view.mock_multi_option) 481 482 def test_extra_package(self): 483 options = PipelineOptions([ 484 '--extra_package', 485 'abc', 486 '--extra_packages', 487 'def', 488 '--extra_packages', 489 'ghi' 490 ]) 491 self.assertEqual( 492 sorted(options.get_all_options()['extra_packages']), 493 ['abc', 'def', 'ghi']) 494 495 options = PipelineOptions(flags=['']) 496 self.assertEqual(options.get_all_options()['extra_packages'], None) 497 498 def test_dataflow_job_file(self): 499 options = PipelineOptions(['--dataflow_job_file', 'abc']) 500 self.assertEqual(options.get_all_options()['dataflow_job_file'], 'abc') 501 502 options = PipelineOptions(flags=['']) 503 self.assertEqual(options.get_all_options()['dataflow_job_file'], None) 504 505 def test_template_location(self): 506 options = PipelineOptions(['--template_location', 'abc']) 507 self.assertEqual(options.get_all_options()['template_location'], 'abc') 508 509 options = PipelineOptions(flags=['']) 510 self.assertEqual(options.get_all_options()['template_location'], None) 511 512 def test_redefine_options(self): 513 class TestRedefinedOptions(PipelineOptions): # pylint: disable=unused-variable 514 @classmethod 515 def _add_argparse_args(cls, parser): 516 parser.add_argument('--redefined_flag', action='store_true') 517 518 class TestRedefinedOptions(PipelineOptions): # pylint: disable=function-redefined 519 @classmethod 520 def _add_argparse_args(cls, parser): 521 parser.add_argument('--redefined_flag', action='store_true') 522 523 options = PipelineOptions(['--redefined_flag']) 524 self.assertTrue(options.get_all_options()['redefined_flag']) 525 526 # TODO(https://github.com/apache/beam/issues/18197): Require unique names 527 # only within a test. For now, <file name acronym>_vp_arg<number> will be 528 # the convention to name value-provider arguments in tests, as opposed to 529 # <file name acronym>_non_vp_arg<number> for non-value-provider arguments. 530 # The number will grow per file as tests are added. 531 def test_value_provider_options(self): 532 class UserOptions(PipelineOptions): 533 @classmethod 534 def _add_argparse_args(cls, parser): 535 parser.add_value_provider_argument( 536 '--pot_vp_arg1', help='This flag is a value provider') 537 538 parser.add_value_provider_argument('--pot_vp_arg2', default=1, type=int) 539 parser.add_argument('--pot_non_vp_arg1', default=1, type=int) 540 541 # Provide values: if not provided, the option becomes of the type runtime vp 542 options = UserOptions(['--pot_vp_arg1', 'hello']) 543 self.assertIsInstance(options.pot_vp_arg1, StaticValueProvider) 544 self.assertIsInstance(options.pot_vp_arg2, RuntimeValueProvider) 545 self.assertIsInstance(options.pot_non_vp_arg1, int) 546 547 # Values can be overwritten 548 options = UserOptions( 549 pot_vp_arg1=5, 550 pot_vp_arg2=StaticValueProvider(value_type=str, value='bye'), 551 pot_non_vp_arg1=RuntimeValueProvider( 552 option_name='foo', value_type=int, default_value=10)) 553 self.assertEqual(options.pot_vp_arg1, 5) 554 self.assertTrue( 555 options.pot_vp_arg2.is_accessible(), 556 '%s is not accessible' % options.pot_vp_arg2) 557 self.assertEqual(options.pot_vp_arg2.get(), 'bye') 558 self.assertFalse(options.pot_non_vp_arg1.is_accessible()) 559 560 with self.assertRaises(RuntimeError): 561 options.pot_non_vp_arg1.get() 562 563 # Converts extra arguments to list value. 564 def test_extra_args(self): 565 options = PipelineOptions([ 566 '--extra_arg', 567 'val1', 568 '--extra_arg', 569 'val2', 570 '--extra_arg=val3', 571 '--unknown_arg', 572 'val4' 573 ]) 574 575 def add_extra_options(parser): 576 parser.add_argument("--extra_arg", action='append') 577 578 self.assertEqual( 579 options.get_all_options( 580 add_extra_args_fn=add_extra_options)['extra_arg'], 581 ['val1', 'val2', 'val3']) 582 583 # The argparse package by default tries to autocomplete option names. This 584 # results in an "ambiguous option" error from argparse when an unknown option 585 # matching multiple known ones are used. This tests that we suppress this 586 # error. 587 def test_unknown_option_prefix(self): 588 # Test that the "ambiguous option" error is suppressed. 589 options = PipelineOptions(['--profi', 'val1']) 590 options.view_as(ProfilingOptions) 591 592 # Test that valid errors are not suppressed. 593 with self.assertRaises(SystemExit): 594 # Invalid option choice. 595 options = PipelineOptions(['--type_check_strictness', 'blahblah']) 596 options.view_as(TypeOptions) 597 598 def test_add_experiment(self): 599 options = PipelineOptions([]) 600 options.view_as(DebugOptions).add_experiment('new_experiment') 601 self.assertEqual(['new_experiment'], 602 options.view_as(DebugOptions).experiments) 603 604 def test_add_experiment_preserves_existing_experiments(self): 605 options = PipelineOptions(['--experiment=existing_experiment']) 606 options.view_as(DebugOptions).add_experiment('new_experiment') 607 self.assertEqual(['existing_experiment', 'new_experiment'], 608 options.view_as(DebugOptions).experiments) 609 610 def test_lookup_experiments(self): 611 options = PipelineOptions([ 612 '--experiment=existing_experiment', 613 '--experiment', 614 'key=value', 615 '--experiment', 616 'master_key=k1=v1,k2=v2', 617 ]) 618 debug_options = options.view_as(DebugOptions) 619 self.assertEqual( 620 'default_value', 621 debug_options.lookup_experiment('nonexistent', 'default_value')) 622 self.assertEqual( 623 'value', debug_options.lookup_experiment('key', 'default_value')) 624 self.assertEqual( 625 'k1=v1,k2=v2', debug_options.lookup_experiment('master_key')) 626 self.assertEqual( 627 True, debug_options.lookup_experiment('existing_experiment')) 628 629 def test_transform_name_mapping(self): 630 options = PipelineOptions(['--transform_name_mapping={\"from\":\"to\"}']) 631 mapping = options.view_as(GoogleCloudOptions).transform_name_mapping 632 self.assertEqual(mapping['from'], 'to') 633 634 def test_dataflow_service_options(self): 635 options = PipelineOptions([ 636 '--dataflow_service_option', 637 'whizz=bang', 638 '--dataflow_service_option', 639 'beep=boop' 640 ]) 641 self.assertEqual( 642 sorted(options.get_all_options()['dataflow_service_options']), 643 ['beep=boop', 'whizz=bang']) 644 645 options = PipelineOptions([ 646 '--dataflow_service_options', 647 'whizz=bang', 648 '--dataflow_service_options', 649 'beep=boop' 650 ]) 651 self.assertEqual( 652 sorted(options.get_all_options()['dataflow_service_options']), 653 ['beep=boop', 'whizz=bang']) 654 655 options = PipelineOptions(flags=['']) 656 self.assertEqual( 657 options.get_all_options()['dataflow_service_options'], None) 658 659 def test_options_store_false_with_different_dest(self): 660 parser = _BeamArgumentParser() 661 for cls in PipelineOptions.__subclasses__(): 662 cls._add_argparse_args(parser) 663 664 actions = parser._actions.copy() 665 options_to_flags = {} 666 options_diff_dest_store_true = {} 667 668 for i in range(len(actions)): 669 flag_names = actions[i].option_strings 670 option_name = actions[i].dest 671 672 if isinstance(actions[i].const, bool): 673 for flag_name in flag_names: 674 flag_name = flag_name.strip('-') 675 if flag_name != option_name: 676 # Capture flags which has store_action=True and has a 677 # different dest. This behavior would be confusing. 678 if actions[i].const: 679 options_diff_dest_store_true[flag_name] = option_name 680 continue 681 # check the flags like no_use_public_ips 682 # default is None, action is {True, False} 683 if actions[i].default is None: 684 options_to_flags[option_name] = flag_name 685 686 self.assertEqual( 687 len(options_diff_dest_store_true), 688 0, 689 _LOGGER.error( 690 "There should be no flags that have a dest " 691 "different from flag name and action as " 692 "store_true. It would be confusing " 693 "to the user. Please specify the dest as the " 694 "flag_name instead.")) 695 from apache_beam.options.pipeline_options import ( 696 _FLAG_THAT_SETS_FALSE_VALUE) 697 698 self.assertDictEqual( 699 _FLAG_THAT_SETS_FALSE_VALUE, 700 options_to_flags, 701 "If you are adding a new boolean flag with default=None," 702 " with different dest/option_name from the flag name, please add " 703 "the dest and the flag name to the map " 704 "_FLAG_THAT_SETS_FALSE_VALUE in PipelineOptions.py") 705 706 707 if __name__ == '__main__': 708 logging.getLogger().setLevel(logging.INFO) 709 unittest.main()