github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/transforms/environments_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  """Unit tests for the transform.environments classes."""
    20  
    21  # pytype: skip-file
    22  
    23  import logging
    24  import unittest
    25  
    26  from apache_beam.options.pipeline_options import PortableOptions
    27  from apache_beam.portability import common_urns
    28  from apache_beam.runners import pipeline_context
    29  from apache_beam.transforms import environments
    30  from apache_beam.transforms.environments import DockerEnvironment
    31  from apache_beam.transforms.environments import EmbeddedPythonEnvironment
    32  from apache_beam.transforms.environments import EmbeddedPythonGrpcEnvironment
    33  from apache_beam.transforms.environments import Environment
    34  from apache_beam.transforms.environments import ExternalEnvironment
    35  from apache_beam.transforms.environments import ProcessEnvironment
    36  from apache_beam.transforms.environments import SubprocessSDKEnvironment
    37  
    38  
    39  class RunnerApiTest(unittest.TestCase):
    40    def test_environment_encoding(self):
    41      for environment in (DockerEnvironment(),
    42                          DockerEnvironment(container_image='img'),
    43                          DockerEnvironment(capabilities=['x, y, z']),
    44                          ProcessEnvironment('run.sh'),
    45                          ProcessEnvironment('run.sh',
    46                                             os='linux',
    47                                             arch='amd64',
    48                                             env={'k1': 'v1'}),
    49                          ExternalEnvironment('localhost:8080'),
    50                          ExternalEnvironment('localhost:8080',
    51                                              params={'k1': 'v1'}),
    52                          EmbeddedPythonEnvironment(),
    53                          EmbeddedPythonGrpcEnvironment(),
    54                          EmbeddedPythonGrpcEnvironment(
    55                              state_cache_size=0, data_buffer_time_limit_ms=0),
    56                          SubprocessSDKEnvironment(command_string=u'foö')):
    57        context = pipeline_context.PipelineContext()
    58        proto = environment.to_runner_api(context)
    59        reconstructed = Environment.from_runner_api(proto, context)
    60        self.assertEqual(environment, reconstructed)
    61        self.assertEqual(proto, reconstructed.to_runner_api(context))
    62  
    63    def test_sdk_capabilities(self):
    64      sdk_capabilities = environments.python_sdk_capabilities()
    65      self.assertIn(common_urns.coders.LENGTH_PREFIX.urn, sdk_capabilities)
    66      self.assertIn(
    67          common_urns.protocols.HARNESS_MONITORING_INFOS.urn, sdk_capabilities)
    68      self.assertIn(common_urns.protocols.WORKER_STATUS.urn, sdk_capabilities)
    69      self.assertIn(
    70          common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn,
    71          sdk_capabilities)
    72      self.assertIn(common_urns.primitives.TO_STRING.urn, sdk_capabilities)
    73  
    74    def test_default_capabilities(self):
    75      environment = DockerEnvironment.from_options(
    76          PortableOptions(sdk_location='container'))
    77      context = pipeline_context.PipelineContext()
    78      proto = environment.to_runner_api(context)
    79      self.assertEqual(
    80          set(proto.capabilities),
    81          set(environments.python_sdk_docker_capabilities()))
    82  
    83  
    84  class EnvironmentOptionsTest(unittest.TestCase):
    85    def test_process_variables_empty(self):
    86      options = PortableOptions([
    87          '--environment_type=PROCESS',
    88          '--environment_option=process_command=foo',
    89          '--sdk_location=container'
    90      ])
    91      environment = ProcessEnvironment.from_options(options)
    92      self.assertEqual(environment.command, 'foo')
    93      self.assertEqual(environment.env, {})
    94  
    95    def test_process_variables_set(self):
    96      options = PortableOptions([
    97          '--environment_type=PROCESS',
    98          '--environment_option=process_command=foo',
    99          '--environment_option=process_variables='
   100          'BASH_VARIABLE_ONE=spam,BASH_VARIABLE_TWO=ham',
   101          '--sdk_location=container'
   102      ])
   103      environment = ProcessEnvironment.from_options(options)
   104      self.assertEqual(environment.command, 'foo')
   105      self.assertEqual(
   106          environment.env, {
   107              'BASH_VARIABLE_ONE': 'spam', 'BASH_VARIABLE_TWO': 'ham'
   108          })
   109  
   110    def test_process_variables_missing_rvalue(self):
   111      with self.assertRaises(ValueError):
   112        options = PortableOptions([
   113            '--environment_type=PROCESS',
   114            '--environment_option=process_command=foo',
   115            '--environment_option=process_variables='
   116            'BASH_VARIABLE_ONE=spam,MISSING_RVALUE',
   117            '--sdk_location=container'
   118        ])
   119        ProcessEnvironment.from_options(options)
   120  
   121    def test_environments_with_same_hints_are_equal(self):
   122      options = PortableOptions([
   123          '--environment_type=PROCESS',
   124          '--environment_option=process_command=foo',
   125          '--sdk_location=container',
   126          '--resource_hint=accelerator=gpu',
   127      ])
   128      environment1 = ProcessEnvironment.from_options(options)
   129      environment2 = ProcessEnvironment.from_options(options)
   130      self.assertEqual(environment1, environment2)
   131  
   132  
   133  if __name__ == '__main__':
   134    logging.getLogger().setLevel(logging.INFO)
   135    unittest.main()