| # | |
| # 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 datetime | |
| import glob | |
| import logging | |
| import os | |
| import tempfile | |
| import unittest | |
| from apache_beam.yaml import main | |
| try: | |
| import jsonschema | |
| except ImportError: | |
| jsonschema = None | |
| TEST_PIPELINE = ''' | |
| pipeline: | |
| type: chain | |
| transforms: | |
| - type: Create | |
| config: | |
| elements: [ELEMENT] | |
| # Writing to an actual file here rather than just using AssertThat | |
| # because this is an integration test and above all we want to ensure | |
| # the pipeline actually runs (and asserts may not fail if there's a | |
| # bug in the invocation logic). | |
| - type: WriteToText | |
| config: | |
| path: PATH | |
| tests: | |
| - name: InlineTest | |
| mock_outputs: | |
| - name: Create | |
| elements: ['a', 'b', 'c'] | |
| expected_inputs: | |
| - name: WriteToText | |
| elements: | |
| - {element: a} | |
| - {element: b} | |
| - {element: c} | |
| ''' | |
| PASSING_TEST_SUITE = ''' | |
| tests: | |
| - name: ExternalTest # comment | |
| mock_outputs: | |
| - name: Create | |
| elements: ['a', 'b', 'c'] | |
| expected_inputs: | |
| - name: WriteToText | |
| elements: | |
| - element: a | |
| - element: b | |
| - element: c | |
| ''' | |
| FAILING_TEST_SUITE = ''' | |
| tests: | |
| - name: ExternalTest # comment | |
| mock_outputs: | |
| - name: Create | |
| elements: ['a', 'b', 'c'] | |
| expected_inputs: | |
| - name: WriteToText | |
| elements: | |
| - element: x | |
| ''' | |
| @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") | |
| class MainTest(unittest.TestCase): | |
| def test_pipeline_spec_from_file(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| yaml_path = os.path.join(tmpdir, 'test.yaml') | |
| out_path = os.path.join(tmpdir, 'out.txt') | |
| with open(yaml_path, 'wt') as fout: | |
| fout.write(TEST_PIPELINE.replace('PATH', out_path)) | |
| main.run(['--yaml_pipeline_file', yaml_path]) | |
| with open(glob.glob(out_path + '*')[0], 'rt') as fin: | |
| self.assertEqual(fin.read().strip(), 'ELEMENT') | |
| def test_pipeline_spec_from_flag(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| out_path = os.path.join(tmpdir, 'out.txt') | |
| main.run(['--yaml_pipeline', TEST_PIPELINE.replace('PATH', out_path)]) | |
| with open(glob.glob(out_path + '*')[0], 'rt') as fin: | |
| self.assertEqual(fin.read().strip(), 'ELEMENT') | |
| def test_jinja_variables(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| out_path = os.path.join(tmpdir, 'out.txt') | |
| main.run([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE.replace('PATH', out_path).replace('ELEMENT', '{{var}}'), | |
| '--jinja_variables', | |
| '{"var": "my_line"}' | |
| ]) | |
| with open(glob.glob(out_path + '*')[0], 'rt') as fin: | |
| self.assertEqual(fin.read().strip(), 'my_line') | |
| def test_jinja_variable_flags(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| out_path = os.path.join(tmpdir, 'out.txt') | |
| main.run([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE.replace('PATH', out_path).replace('ELEMENT', '{{var}}'), | |
| '--jinja_variable_flags=var', | |
| '--var=my_line' | |
| ]) | |
| with open(glob.glob(out_path + '*')[0], 'rt') as fin: | |
| self.assertEqual(fin.read().strip(), 'my_line') | |
| def test_preparse_jinja_flags(self): | |
| argv = [ | |
| '--jinja_variables={"from_vars": 1, "from_both": 2}', | |
| '--jinja_variable_flags=from_both,from_flag,from_missing_flag', | |
| '--from_both=30', | |
| '--from_flag=40', | |
| '--another_arg=foo', | |
| 'pos_arg', | |
| ] | |
| self.assertCountEqual( | |
| main._preparse_jinja_flags(argv), | |
| [ | |
| '--jinja_variables=' | |
| '{"from_vars": 1, "from_both": "30", "from_flag": "40"}', | |
| '--another_arg=foo', | |
| 'pos_arg', | |
| ]) | |
| def test_preparse_jinja_flags_pipeline_option_collision(self): | |
| # A jinja_variable_flags entry that collides with a known pipeline | |
| # option (e.g. runner) must not swallow the pipeline flag. | |
| argv = [ | |
| '--jinja_variable_flags=runner,var', | |
| '--runner=DirectRunner', | |
| '--var=my_line', | |
| ] | |
| self.assertCountEqual( | |
| main._preparse_jinja_flags(argv), [ | |
| '--runner=DirectRunner', | |
| '--jinja_variables=' + '{"var": "my_line"}', | |
| ]) | |
| def test_jinja_datetime(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| out_path = os.path.join(tmpdir, 'out.txt') | |
| main.run([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE.replace('PATH', out_path).replace( | |
| 'ELEMENT', '"{{datetime.datetime.now().strftime("%Y-%m-%d")}}"'), | |
| ]) | |
| with open(glob.glob(out_path + '*')[0], 'rt') as fin: | |
| self.assertEqual( | |
| fin.read().strip(), datetime.datetime.now().strftime("%Y-%m-%d")) | |
| def test_inline_test_specs(self): | |
| main.run_tests(['--yaml_pipeline', TEST_PIPELINE, '--test'], exit=False) | |
| def test_external_test_specs(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| good_suite = os.path.join(tmpdir, 'good.yaml') | |
| with open(good_suite, 'w') as fout: | |
| fout.write(PASSING_TEST_SUITE) | |
| bad_suite = os.path.join(tmpdir, 'bad.yaml') | |
| with open(bad_suite, 'w') as fout: | |
| fout.write(FAILING_TEST_SUITE) | |
| # Must pass. | |
| main.run_tests([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE, | |
| '--test_suite', | |
| good_suite, | |
| ], | |
| exit=False) | |
| # Must fail. (Ensures testing is not a no-op.) | |
| with self.assertRaisesRegex(Exception, 'errors=1 failures=0'): | |
| main.run_tests([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE, | |
| '--test_suite', | |
| bad_suite, | |
| ], | |
| exit=False) | |
| def test_fix_suite(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| test_suite = os.path.join(tmpdir, 'tests.yaml') | |
| with open(test_suite, 'w') as fout: | |
| fout.write(FAILING_TEST_SUITE) | |
| main.run_tests([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE, | |
| '--test_suite', | |
| test_suite, | |
| '--fix_tests' | |
| ], | |
| exit=False) | |
| with open(test_suite) as fin: | |
| self.assertEqual(fin.read(), PASSING_TEST_SUITE) | |
| def test_create_test(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| test_suite = os.path.join(tmpdir, 'tests.yaml') | |
| with open(test_suite, 'w') as fout: | |
| fout.write('') | |
| main.run_tests([ | |
| '--yaml_pipeline', | |
| TEST_PIPELINE.replace('ELEMENT', 'x'), | |
| '--test_suite', | |
| test_suite, | |
| '--create_test' | |
| ], | |
| exit=False) | |
| with open(test_suite) as fin: | |
| self.assertEqual( | |
| fin.read(), | |
| ''' | |
| tests: | |
| - mock_outputs: [] | |
| expected_inputs: | |
| - name: WriteToText | |
| elements: | |
| - element: x | |
| '''.lstrip()) | |
| if __name__ == '__main__': | |
| logging.getLogger().setLevel(logging.INFO) | |
| unittest.main() |