github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/internal/cloudpickle_pickler_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Unit tests for the cloudpickle_pickler module."""
    19  
    20  # pytype: skip-file
    21  
    22  import sys
    23  import threading
    24  import types
    25  import unittest
    26  
    27  from apache_beam.internal import module_test
    28  from apache_beam.internal.cloudpickle_pickler import dumps
    29  from apache_beam.internal.cloudpickle_pickler import loads
    30  
    31  
    32  class PicklerTest(unittest.TestCase):
    33  
    34    NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
    35  
    36    def test_basics(self):
    37      self.assertEqual([1, 'a', (u'z', )], loads(dumps([1, 'a', (u'z', )])))
    38      fun = lambda x: 'xyz-%s' % x
    39      self.assertEqual('xyz-abc', loads(dumps(fun))('abc'))
    40  
    41    def test_lambda_with_globals(self):
    42      """Tests that the globals of a function are preserved."""
    43  
    44      # The point of the test is that the lambda being called after unpickling
    45      # relies on having the re module being loaded.
    46      self.assertEqual(['abc', 'def'],
    47                       loads(dumps(
    48                           module_test.get_lambda_with_globals()))('abc def'))
    49  
    50    def test_lambda_with_main_globals(self):
    51      self.assertEqual(unittest, loads(dumps(lambda: unittest))())
    52  
    53    def test_lambda_with_closure(self):
    54      """Tests that the closure of a function is preserved."""
    55      self.assertEqual(
    56          'closure: abc',
    57          loads(dumps(module_test.get_lambda_with_closure('abc')))())
    58  
    59    def test_class_object_pickled(self):
    60      self.assertEqual(['abc', 'def'],
    61                       loads(dumps(module_test.Xyz))().foo('abc def'))
    62  
    63    def test_class_instance_pickled(self):
    64      self.assertEqual(['abc', 'def'],
    65                       loads(dumps(module_test.XYZ_OBJECT)).foo('abc def'))
    66  
    67    def test_pickling_preserves_closure_of_a_function(self):
    68      self.assertEqual(
    69          'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum)
    70      self.assertEqual(
    71          'Y:abc',
    72          loads(dumps(module_test.TopClass.MiddleClass.NestedClass('abc'))).datum)
    73  
    74    def test_pickle_dynamic_class(self):
    75      self.assertEqual(
    76          'Z:abc', loads(dumps(module_test.create_class('abc'))).get())
    77  
    78    def test_generators(self):
    79      with self.assertRaises(TypeError):
    80        dumps((_ for _ in range(10)))
    81  
    82    def test_recursive_class(self):
    83      self.assertEqual(
    84          'RecursiveClass:abc',
    85          loads(dumps(module_test.RecursiveClass('abc').datum)))
    86  
    87    def test_function_with_external_reference(self):
    88      out_of_scope_var = 'expected_value'
    89  
    90      def foo():
    91        return out_of_scope_var
    92  
    93      self.assertEqual('expected_value', loads(dumps(foo))())
    94  
    95    def test_pickle_rlock(self):
    96      rlock_instance = threading.RLock()
    97      rlock_type = type(rlock_instance)
    98  
    99      self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type)
   100  
   101    @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced')
   102    def test_dump_and_load_mapping_proxy(self):
   103      self.assertEqual(
   104          'def', loads(dumps(types.MappingProxyType({'abc': 'def'})))['abc'])
   105      self.assertEqual(
   106          types.MappingProxyType, type(loads(dumps(types.MappingProxyType({})))))
   107  
   108    # pylint: disable=exec-used
   109    @unittest.skipIf(sys.version_info < (3, 7), 'Python 3.7 or above only')
   110    def test_dataclass(self):
   111      exec(
   112          '''
   113  from apache_beam.internal.module_test import DataClass
   114  self.assertEqual(DataClass(datum='abc'), loads(dumps(DataClass(datum='abc'))))
   115      ''')
   116  
   117  
   118  if __name__ == '__main__':
   119    unittest.main()