| # |
| # 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. |
| # |
| import logging |
| import unittest |
| |
| import yaml |
| |
| import apache_beam as beam |
| from apache_beam.yaml import YamlTransform |
| from apache_beam.yaml import yaml_provider |
| from apache_beam.yaml.yaml_provider import InlineProvider |
| from apache_beam.yaml.yaml_transform import SafeLineLoader |
| from apache_beam.yaml.yaml_transform import Scope |
| from apache_beam.yaml.yaml_transform import chain_as_composite |
| from apache_beam.yaml.yaml_transform import ensure_errors_consumed |
| from apache_beam.yaml.yaml_transform import ensure_transforms_have_types |
| from apache_beam.yaml.yaml_transform import expand_composite_transform |
| from apache_beam.yaml.yaml_transform import extract_name |
| from apache_beam.yaml.yaml_transform import identify_object |
| from apache_beam.yaml.yaml_transform import normalize_inputs_outputs |
| from apache_beam.yaml.yaml_transform import normalize_source_sink |
| from apache_beam.yaml.yaml_transform import only_element |
| from apache_beam.yaml.yaml_transform import pipeline_as_composite |
| from apache_beam.yaml.yaml_transform import preprocess |
| from apache_beam.yaml.yaml_transform import preprocess_flattened_inputs |
| from apache_beam.yaml.yaml_transform import preprocess_windowing |
| from apache_beam.yaml.yaml_transform import push_windowing_to_roots |
| |
| |
| class SafeLineLoaderTest(unittest.TestCase): |
| def test_get_line(self): |
| pipeline_yaml = ''' |
| type: composite |
| input: |
| elements: input |
| transforms: |
| - type: PyMap |
| name: Square |
| input: elements |
| config: |
| fn: "lambda x: x * x" |
| - type: PyMap |
| name: Cube |
| input: elements |
| config: |
| fn: "lambda x: x * x * x" |
| output: |
| Flatten |
| ''' |
| spec = yaml.load(pipeline_yaml, Loader=SafeLineLoader) |
| self.assertEqual(SafeLineLoader.get_line(spec['type']), 2) |
| self.assertEqual(SafeLineLoader.get_line(spec['input']), 4) |
| self.assertEqual(SafeLineLoader.get_line(spec['transforms'][0]), 6) |
| self.assertEqual(SafeLineLoader.get_line(spec['transforms'][0]['type']), 6) |
| self.assertEqual(SafeLineLoader.get_line(spec['transforms'][0]['name']), 7) |
| self.assertEqual(SafeLineLoader.get_line(spec['transforms'][1]), 11) |
| self.assertEqual(SafeLineLoader.get_line(spec['output']), 17) |
| self.assertEqual(SafeLineLoader.get_line(spec['transforms']), "unknown") |
| |
| def test_strip_metadata(self): |
| spec_yaml = ''' |
| transforms: |
| - type: PyMap |
| name: Square |
| ''' |
| spec = yaml.load(spec_yaml, Loader=SafeLineLoader) |
| stripped = SafeLineLoader.strip_metadata(spec['transforms']) |
| |
| self.assertFalse(hasattr(stripped[0], '__line__')) |
| self.assertFalse(hasattr(stripped[0], '__uuid__')) |
| |
| def test_strip_metadata_nothing_to_strip(self): |
| spec_yaml = 'prop: 123' |
| spec = yaml.load(spec_yaml, Loader=SafeLineLoader) |
| stripped = SafeLineLoader.strip_metadata(spec['prop']) |
| |
| self.assertFalse(hasattr(stripped, '__line__')) |
| self.assertFalse(hasattr(stripped, '__uuid__')) |
| |
| |
| def new_pipeline(): |
| return beam.Pipeline( |
| options=beam.options.pipeline_options.PipelineOptions( |
| pickle_library='cloudpickle')) |
| |
| |
| class MainTest(unittest.TestCase): |
| def assertYaml(self, expected, result): |
| result = SafeLineLoader.strip_metadata(result) |
| expected = yaml.load(expected, Loader=SafeLineLoader) |
| expected = SafeLineLoader.strip_metadata(expected) |
| |
| self.assertEqual(expected, result) |
| |
| def get_scope_by_spec(self, p, spec, inputs=None): |
| if inputs is None: |
| inputs = {} |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| |
| scope = Scope( |
| beam.pvalue.PBegin(p), |
| inputs, |
| spec['transforms'], |
| yaml_provider.standard_providers(), {}) |
| return scope, spec |
| |
| def test_pipeline_as_composite_with_type_transforms(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = pipeline_as_composite(spec) |
| |
| self.assertEqual(result['type'], 'composite') |
| self.assertEqual(result['name'], None) |
| |
| def test_pipeline_as_composite_with_transforms(self): |
| spec = ''' |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = pipeline_as_composite(spec) |
| |
| self.assertEqual(result['type'], 'composite') |
| self.assertEqual(result['name'], None) |
| |
| def test_pipeline_as_composite_list(self): |
| spec = ''' |
| - type: Create |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = pipeline_as_composite(spec) |
| |
| expected = ''' |
| type: composite |
| name: null |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_expand_composite_transform_with_name(self): |
| with new_pipeline() as p: |
| spec = ''' |
| type: composite |
| name: Custom |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| output: |
| Create |
| ''' |
| scope, spec = self.get_scope_by_spec(p, spec) |
| self.assertRegex( |
| str(expand_composite_transform(spec, scope)['output']), |
| r"PCollection.*Custom/Create/Map.*") |
| |
| def test_expand_composite_transform_with_name_input(self): |
| with new_pipeline() as p: |
| spec = ''' |
| type: composite |
| input: elements |
| transforms: |
| - type: LogForTesting |
| input: input |
| output: |
| LogForTesting |
| ''' |
| elements = p | beam.Create(range(3)) |
| scope, spec = self.get_scope_by_spec(p, spec, |
| inputs={'elements': elements}) |
| self.assertRegex( |
| str(expand_composite_transform(spec, scope)['output']), |
| r"PCollection.*Composite/Map.*") |
| |
| def test_expand_composite_transform_root(self): |
| with new_pipeline() as p: |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| output: |
| Create |
| ''' |
| scope, spec = self.get_scope_by_spec(p, spec) |
| self.assertRegex( |
| str(expand_composite_transform(spec, scope)['output']), |
| r"PCollection.*Create/Map.*") |
| |
| def test_chain_as_composite(self): |
| spec = ''' |
| type: chain |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = chain_as_composite(spec) |
| |
| expected = f''' |
| type: composite |
| name: Chain |
| transforms: |
| - type: Create |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| input: {spec['transforms'][0]['__uuid__']} |
| output: {spec['transforms'][1]['__uuid__']} |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_chain_as_composite_with_wrong_output_type(self): |
| spec = ''' |
| type: chain |
| transforms: |
| - type: Create |
| elements: [0,1,2] |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| output: |
| Create |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| with self.assertRaisesRegex(ValueError, |
| r"Explicit output.*of the chain transform is " |
| r"not an output of the last transform"): |
| chain_as_composite(spec) |
| |
| def test_chain_as_composite_with_wrong_output_name(self): |
| spec = ''' |
| type: chain |
| transforms: |
| - type: Create |
| name: elements |
| elements: [0,1,2] |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| output: |
| elements |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| with self.assertRaisesRegex(ValueError, |
| r"Explicit output.*of the chain transform is " |
| r"not an output of the last transform"): |
| chain_as_composite(spec) |
| |
| def test_chain_as_composite_with_outputs_override(self): |
| spec = ''' |
| type: chain |
| transforms: |
| - type: Create |
| elements: [0,1,2] |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| output: |
| PyMap |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual( |
| chain_as_composite(spec)['output']['output'], |
| f"{spec['transforms'][1]['__uuid__']}.PyMap") |
| |
| def test_chain_as_composite_with_input(self): |
| spec = ''' |
| type: chain |
| input: |
| elements |
| transforms: |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual( |
| chain_as_composite(spec)['transforms'][0]['input'], {"input": "input"}) |
| |
| def test_normalize_source_sink(self): |
| spec = ''' |
| source: |
| type: Create |
| config: |
| elements: [0,1,2] |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| sink: |
| type: PyMap |
| config: |
| fn: "lambda x: x + 41" |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_source_sink(spec) |
| |
| expected = ''' |
| transforms: |
| - type: Create |
| input: {'__explicitly_empty__': null} |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| - type: PyMap |
| config: |
| fn: "lambda x: x + 41" |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_normalize_source_sink_only_source(self): |
| spec = ''' |
| source: |
| type: Create |
| config: |
| elements: [0,1,2] |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_source_sink(spec) |
| |
| expected = ''' |
| transforms: |
| - type: Create |
| input: {'__explicitly_empty__': null} |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_normalize_source_sink_only_sink(self): |
| spec = ''' |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| sink: |
| type: PyMap |
| config: |
| fn: "lambda x: x + 41" |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_source_sink(spec) |
| expected = ''' |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| - type: PyMap |
| config: |
| fn: "lambda x: x + 41" |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_normalize_source_sink_no_source_no_sink(self): |
| spec = ''' |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_source_sink(spec) |
| |
| expected = ''' |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_source_sink_composite(self): |
| spec = ''' |
| type: composite |
| source: |
| type: Create |
| config: |
| elements: [0,1,2] |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_source_sink(spec) |
| |
| expected = ''' |
| type: composite |
| transforms: |
| - type: Create |
| input: {'__explicitly_empty__': null} |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_source_sink_chain(self): |
| spec = ''' |
| type: chain |
| source: |
| type: Create |
| config: |
| elements: [0,1,2] |
| transforms: |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_source_sink(spec) |
| |
| expected = ''' |
| type: chain |
| transforms: |
| - type: Create |
| input: {'__explicitly_empty__': null} |
| config: |
| elements: [0,1,2] |
| - type: PyMap |
| config: |
| fn: 'lambda x: x*x' |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_source_sink_other(self): |
| spec = ''' |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual(normalize_source_sink(spec), spec) |
| |
| def test_normalize_inputs_outputs(self): |
| spec = ''' |
| type: PyMap |
| input: [Create1, Create2] |
| fn: 'lambda x: x*x' |
| output: Squared |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_inputs_outputs(spec) |
| |
| expected = ''' |
| type: PyMap |
| input: |
| input: [Create1, Create2] |
| fn: 'lambda x: x*x' |
| output: |
| output: Squared |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_normalize_inputs_outputs_dict(self): |
| spec = ''' |
| type: PyMap |
| input: [Create1, Create2] |
| fn: 'lambda x: x*x' |
| output: |
| out1: Squared1 |
| out2: Squared2 |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = normalize_inputs_outputs(spec) |
| |
| expected = ''' |
| type: PyMap |
| input: |
| input: [Create1, Create2] |
| fn: 'lambda x: x*x' |
| output: |
| out1: Squared1 |
| out2: Squared2 |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_identify_object_with_name(self): |
| spec = ''' |
| type: PyMap |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertRegex(identify_object(spec), r"PyMap.*[0-9]") |
| |
| def test_identify_object(self): |
| spec = ''' |
| argument: PyMap |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertRegex(identify_object(spec), r"at.*[0-9]") |
| |
| def test_extract_name_by_type(self): |
| spec = ''' |
| type: PyMap |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual(extract_name(spec), "PyMap") |
| |
| def test_extract_name_by_id(self): |
| spec = ''' |
| type: PyMap |
| id: PyMapId |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual(extract_name(spec), "PyMapId") |
| |
| def test_extract_name_by_name(self): |
| spec = ''' |
| type: PyMap |
| name: PyMapName |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual(extract_name(spec), "PyMapName") |
| |
| def test_extract_name_no_name(self): |
| spec = ''' |
| transforms: |
| - arg: PyMap |
| fn: 'lambda x: x*x' |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| self.assertEqual(extract_name(spec), "") |
| |
| def test_push_windowing_to_roots(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Create |
| elements: [0,1,2] |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| input: Create |
| windowing: |
| type: fixed |
| size: 2 |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| |
| result = push_windowing_to_roots(spec) |
| |
| expected = ''' |
| type: composite |
| transforms: |
| - type: Create |
| elements: [0,1,2] |
| windowing: |
| type: fixed |
| size: 2 |
| __consumed_outputs: |
| - null |
| input: {} |
| output: {} |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| input: |
| input: Create |
| output: {} |
| windowing: |
| type: fixed |
| size: 2 |
| input: {} |
| output: {} |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_windowing_custom_type(self): |
| spec = ''' |
| type: SumGlobally |
| input: Create |
| windowing: |
| type: fixed |
| size: 4 |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| result = preprocess_windowing(spec) |
| |
| expected = f''' |
| type: composite |
| name: SumGlobally |
| input: |
| input: Create |
| transforms: |
| - type: SumGlobally |
| input: |
| input: {result['transforms'][1]['__uuid__']} |
| output: {{}} |
| - type: WindowInto |
| name: WindowInto[input] |
| windowing: |
| type: fixed |
| size: 4 |
| input: input |
| output: {result['transforms'][0]['__uuid__']} |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_windowing_composite_with_windowing_outer(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: CreateTimestamped |
| name: Create |
| elements: [0, 2, 4] |
| - type: SumGlobally |
| input: Create |
| windowing: |
| type: fixed |
| size: 4 |
| output: SumGlobally |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| |
| result = preprocess_windowing(spec) |
| |
| expected = ''' |
| type: composite |
| input: {} |
| transforms: |
| - type: CreateTimestamped |
| name: Create |
| elements: [0, 2, 4] |
| windowing: |
| type: fixed |
| size: 4 |
| __consumed_outputs: |
| - null |
| input: {} |
| output: {} |
| - type: SumGlobally |
| input: |
| input: Create |
| output: {} |
| output: |
| output: SumGlobally |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_windowing_composite_with_windowing_on_input(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: CreateTimestamped |
| name: Create |
| elements: [0, 2, 4] |
| - type: SumGlobally |
| input: Create |
| windowing: |
| type: fixed |
| size: 4 |
| output: SumGlobally |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| |
| result = preprocess_windowing(spec) |
| |
| expected = ''' |
| type: composite |
| input: {} |
| transforms: |
| - type: CreateTimestamped |
| name: Create |
| elements: [0, 2, 4] |
| input: {} |
| output: {} |
| - type: SumGlobally |
| input: |
| input: Create |
| windowing: |
| type: fixed |
| size: 4 |
| output: {} |
| output: |
| output: SumGlobally |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_windowing_other_type_with_no_inputs(self): |
| spec = ''' |
| type: CreateTimestamped |
| name: Create |
| elements: [0, 2, 4] |
| windowing: |
| type: fixed |
| size: 4 |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| result = preprocess_windowing(spec) |
| |
| expected = f''' |
| type: composite |
| name: Create |
| transforms: |
| - type: CreateTimestamped |
| name: Create |
| elements: [0, 2, 4] |
| input: {{}} |
| output: {{}} |
| - type: WindowInto |
| name: WindowInto[None] |
| input: |
| input: {result['transforms'][0]["__uuid__"]} |
| windowing: |
| type: fixed |
| size: 4 |
| output: {result['transforms'][1]["__uuid__"]} |
| ''' |
| self.maxDiff = 1e9 |
| |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_flattened_inputs_implicit(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| input: [Create1, Create2] |
| output: CreateTimestamped |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| result = preprocess_flattened_inputs(spec) |
| |
| expected = f''' |
| type: composite |
| transforms: |
| - type: Flatten |
| name: PyMap-Flatten[input] |
| input: |
| input0: Create1 |
| input1: Create2 |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| input: |
| input: {result['transforms'][0]['__uuid__']} |
| output: {{}} |
| output: CreateTimestamped |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_preprocess_flattened_inputs_explicit_flatten(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Flatten |
| input: [Create1, Create2] |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| input: Flatten |
| output: CreateTimestamped |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| result = preprocess_flattened_inputs(spec) |
| |
| expected = ''' |
| type: composite |
| transforms: |
| - type: Flatten |
| input: |
| input0: Create1 |
| input1: Create2 |
| output: {} |
| - type: PyMap |
| fn: 'lambda x: x*x' |
| input: |
| input: Flatten |
| output: {} |
| output: CreateTimestamped |
| ''' |
| self.assertYaml(expected, result) |
| |
| def test_ensure_transforms_have_types(self): |
| spec = ''' |
| type: PyMap |
| fn: 'lambda x: x*x' |
| input: Flatten |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = ensure_transforms_have_types(spec) |
| self.assertEqual(result, spec) |
| |
| def test_ensure_transforms_have_types_error(self): |
| spec = ''' |
| name: PyMap |
| fn: 'lambda x: x*x' |
| input: Flatten |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| with self.assertRaisesRegex(ValueError, r"Missing type .*"): |
| ensure_transforms_have_types(spec) |
| with self.assertRaisesRegex(ValueError, r"Missing type .*"): |
| preprocess(spec) |
| |
| def test_ensure_transforms_have_providers_error(self): |
| spec = ''' |
| type: UnknownType |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| with self.assertRaisesRegex(ValueError, |
| r"Unknown type or missing provider .*"): |
| preprocess(spec, known_transforms=['KnownType']) |
| |
| def test_ensure_errors_consumed_unconsumed(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Create |
| elements: [1,2,3] |
| - type: MyTransform |
| input: Create |
| error_handling: |
| output: errors |
| output: |
| good: MyTransform |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| with self.assertRaisesRegex(ValueError, r"Unconsumed error.*"): |
| ensure_errors_consumed(spec) |
| |
| def test_ensure_errors_consumed_in_transform(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Create |
| elements: [1,2,3] |
| - type: MyTransform |
| input: Create |
| error_handling: |
| output: errors |
| - name: SaveToFile |
| type: PyMap |
| input: MyTransform.errors |
| output: |
| good: MyTransform |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| result = ensure_errors_consumed(spec) |
| self.assertEqual(result, spec) |
| self.assertEqual(result['transforms'], spec['transforms']) |
| |
| def test_ensure_errors_consumed_no_output_in_error_handling(self): |
| spec = ''' |
| type: composite |
| transforms: |
| - type: Create |
| elements: [1,2,3] |
| - type: MyTransform |
| input: Create |
| error_handling: |
| arg: errors |
| - name: SaveToFile |
| type: PyMap |
| input: MyTransform.errors |
| output: |
| good: MyTransform |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| spec = normalize_inputs_outputs(spec) |
| spec['transforms'] = [ |
| normalize_inputs_outputs(t) for t in spec['transforms'] |
| ] |
| with self.assertRaisesRegex(ValueError, r"Missing output.*"): |
| ensure_errors_consumed(spec) |
| |
| def test_only_element(self): |
| self.assertEqual(only_element((1, )), 1) |
| |
| |
| class YamlTransformTest(unittest.TestCase): |
| def test_init_with_string(self): |
| provider1 = InlineProvider({"MyTransform1": lambda: beam.Map(lambda x: x)}) |
| provider2 = InlineProvider({"MyTransform2": lambda: beam.Map(lambda x: x)}) |
| |
| providers_dict = {"p1": [provider1], "p2": [provider2]} |
| |
| spec = ''' |
| type: chain |
| transforms: |
| - type: Create |
| elements: [1,2,3] |
| - type: LogForTesting |
| ''' |
| result = YamlTransform(spec, providers_dict) |
| self.assertIn('p1', result._providers) # check for custom providers |
| self.assertIn('p2', result._providers) # check for custom providers |
| self.assertIn( |
| 'LogForTesting', result._providers) # check for standard provider |
| self.assertEqual(result._spec['type'], "composite") # preprocessed spec |
| |
| def test_init_with_dict(self): |
| spec = ''' |
| type: chain |
| transforms: |
| - type: Create |
| config: |
| elements: [1,2,3] |
| - type: LogForTesting |
| ''' |
| spec = yaml.load(spec, Loader=SafeLineLoader) |
| result = YamlTransform(spec) |
| self.assertIn( |
| 'LogForTesting', result._providers) # check for standard provider |
| self.assertEqual(result._spec['type'], "composite") # preprocessed spec |
| |
| |
| if __name__ == '__main__': |
| logging.getLogger().setLevel(logging.INFO) |
| unittest.main() |
