github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/examples/snippets/util_test.py (about)

     1  # coding=utf-8
     2  #
     3  # Licensed to the Apache Software Foundation (ASF) under one or more
     4  # contributor license agreements.  See the NOTICE file distributed with
     5  # this work for additional information regarding copyright ownership.
     6  # The ASF licenses this file to You under the Apache License, Version 2.0
     7  # (the "License"); you may not use this file except in compliance with
     8  # the License.  You may obtain a copy of the License at
     9  #
    10  #    http://www.apache.org/licenses/LICENSE-2.0
    11  #
    12  # Unless required by applicable law or agreed to in writing, software
    13  # distributed under the License is distributed on an "AS IS" BASIS,
    14  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    15  # See the License for the specific language governing permissions and
    16  # limitations under the License.
    17  #
    18  
    19  # pytype: skip-file
    20  
    21  import unittest
    22  
    23  from mock import patch
    24  
    25  import apache_beam as beam
    26  from apache_beam.examples.snippets import util
    27  from apache_beam.testing.test_pipeline import TestPipeline
    28  
    29  
    30  class UtilTest(unittest.TestCase):
    31    def test_assert_matches_stdout_object(self):
    32      expected = [
    33          "{'a': '🍓', 'b': True}",
    34          "{'a': '🥕', 'b': 42}",
    35          "{'a': '🍆', 'b': '\"hello\"'}",
    36          "{'a': '🍅', 'b': [1, 2, 3]}",
    37          "{'b': 'B', 'a': '🥔'}",
    38      ]
    39      with TestPipeline() as pipeline:
    40        actual = (
    41            pipeline
    42            | beam.Create([
    43                {
    44                    'a': '🍓', 'b': True
    45                },
    46                {
    47                    'a': '🥕', 'b': 42
    48                },
    49                {
    50                    'a': '🍆', 'b': '"hello"'
    51                },
    52                {
    53                    'a': '🍅', 'b': [1, 2, 3]
    54                },
    55                {
    56                    'a': '🥔', 'b': 'B'
    57                },
    58            ])
    59            | beam.Map(str))
    60        util.assert_matches_stdout(actual, expected)
    61  
    62    def test_assert_matches_stdout_string(self):
    63      expected = ['🍓', '🥕', '🍆', '🍅', '🥔']
    64      with TestPipeline() as pipeline:
    65        actual = (
    66            pipeline
    67            | beam.Create(['🍓', '🥕', '🍆', '🍅', '🥔'])
    68            | beam.Map(str))
    69        util.assert_matches_stdout(actual, expected)
    70  
    71    def test_assert_matches_stdout_sorted_keys(self):
    72      expected = [{'list': [1, 2]}, {'list': [3, 4]}]
    73      with TestPipeline() as pipeline:
    74        actual = (
    75            pipeline
    76            | beam.Create([{
    77                'list': [2, 1]
    78            }, {
    79                'list': [4, 3]
    80            }])
    81            | beam.Map(str))
    82        util.assert_matches_stdout(
    83            actual, expected, lambda elem: {'sorted': sorted(elem['list'])})
    84  
    85    @patch('subprocess.call', lambda cmd: None)
    86    def test_run_shell_commands(self):
    87      commands = [
    88          '  # this is a comment  ',
    89          '  !  echo   this   is   a   shell   command  ',
    90          '  !echo {variable}  ',
    91          '  echo "quoted arguments work"  # trailing comment  ',
    92      ]
    93      actual = list(util.run_shell_commands(commands, variable='hello world'))
    94      expected = [
    95          ['echo', 'this', 'is', 'a', 'shell', 'command'],
    96          ['echo', 'hello', 'world'],
    97          ['echo', 'quoted arguments work'],
    98      ]
    99      self.assertEqual(actual, expected)
   100  
   101  
   102  if __name__ == '__main__':
   103    unittest.main()