jk-shah commented on code in PR #40968: URL: https://github.com/apache/beam/pull/40968#discussion_r4488741706
########## sdks/python/apache_beam/cli/main_test.py: ########## @@ -0,0 +1,390 @@ +# +# 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 contextlib +import glob +import io +import json +import os +import shutil +import tempfile +import unittest +from unittest import mock + +import apache_beam as beam +from apache_beam.cli import main as cli +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.utils import subprocess_server + +try: + import jsonschema + from apache_beam.yaml import main as yaml_main +except ImportError: + jsonschema = None + yaml_main = None + +# Paths are formatted with `json.dumps` so Windows backslashes (such as `\U` +# and `\t` in `C:\Users\...\temp`) are escaped as valid YAML strings. +PIPELINE = ''' +pipeline: + type: chain + transforms: + - type: ReadFromCsv + config: {path: %(input)s} + - type: Filter + config: {language: python, keep: "age > 20"} + - type: WriteToJson + config: {path: %(output)s} +''' + +FAILING_PIPELINE = ''' +pipeline: + type: chain + transforms: + - type: ReadFromCsv + config: {path: %(input)s} + - type: MapToFields + config: {language: python, fields: {x: "age // 0"}} + - type: WriteToJson + config: {path: %(output)s} +''' + +INVALID_PIPELINE = ''' +pipeline: + type: chain + transforms: + - type: NoSuchTransform +''' + + +def _run_cli(argv): + """Runs `beam` with argv; returns (exit code, stdout, stderr).""" + out, err = io.StringIO(), io.StringIO() + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + code = cli.main(argv) + return code, out.getvalue(), err.getvalue() + + +def _read_lines(prefix): + lines = [] + for path in glob.glob(prefix + '*'): + with open(path) as f: + lines.extend(f.read().splitlines()) + return sorted(lines) + + +class ParserTest(unittest.TestCase): + def test_usage_errors_exit_with_2(self): + for argv in ([], ['run'], ['nosuchcommand', 'x']): + with self.subTest(argv=argv): + with contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit) as e: + cli.main(argv) + self.assertEqual(e.exception.code, cli.EXIT_USAGE) + + def test_missing_file(self): + code, _, err = _run_cli(['run', 'no/such/pipeline.yaml']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn("no such file: 'no/such/pipeline.yaml'", err) + self.assertNotIn('pipeline options go after FILE', err) + + def test_missing_file_hints_at_option_order(self): + # argparse takes "dataflow" for FILE when an option uses a space. + code, _, err = _run_cli(['run', '--runner', 'dataflow', 'p.yaml']) + self.assertEqual(code, cli.EXIT_USAGE) + self.assertIn("no such file: 'dataflow'", err) + self.assertIn('pipeline options go after FILE', err) + + [email protected]( + yaml_main is None or jsonschema is None, 'requires the [yaml] extra') +class RunTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp, ignore_errors=True) + self.input = os.path.join(self.tmp, 'users.csv') + with open(self.input, 'w') as f: + f.write('name,age\nAda,36\nLinus,17\nGrace,45\n') + + def _pipeline_file(self, template, output='users.json', input_path=None): + path = os.path.join(self.tmp, 'pipeline.yaml') + with open(path, 'w') as f: + f.write( + template % { + 'input': json.dumps(input_path or self.input), + 'output': json.dumps(os.path.join(self.tmp, output)), + }) + return path + + def test_runs_like_yaml_main(self): + # The CLI and `python -m apache_beam.yaml.main` write the same output for + # the same pipeline file and options. + path = self._pipeline_file(PIPELINE, 'cli_users.json') + code, out, _ = _run_cli(['run', path, '--runner=DirectRunner']) + self.assertEqual(code, cli.EXIT_OK) + self.assertRegex(out, r'Done\. Ran locally on (Prism|the Python direct)') + + reference = self._pipeline_file(PIPELINE, 'reference_users.json') + yaml_main.run( + ['--yaml_pipeline_file=' + reference, '--runner=DirectRunner']) + cli_rows = _read_lines(os.path.join(self.tmp, 'cli_users.json')) + self.assertEqual(len(cli_rows), 2) + self.assertEqual( + cli_rows, _read_lines(os.path.join(self.tmp, 'reference_users.json'))) + + def test_backslash_sequences_in_paths(self): + # Regression test for Windows temp paths (`C:\Users\...\tmp...`), where + # unescaped `\U` and `\t` inside double-quoted YAML strings fail parsing. + tricky_input = os.path.join(self.tmp, r'C_\Users\tmp_users.csv') + shutil.copyfile(self.input, tricky_input) Review Comment: Thanks @damccorm! Fixed in 153569ffcb4 by creating os.path.dirname(tricky_input) before shutil.copyfile. I am good on following up with a website documentation PR (and a Beam blog post!) after the 2.79 release — happy to put those together once the subcommands land. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
