# # 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_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()