github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/internal/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 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.pickler import dumps
    29  from apache_beam.internal.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(self):
    60      """Tests that a class object is pickled correctly."""
    61      self.assertEqual(['abc', 'def'],
    62                       loads(dumps(module_test.Xyz))().foo('abc def'))
    63  
    64    def test_object(self):
    65      """Tests that a class instance is pickled correctly."""
    66      self.assertEqual(['abc', 'def'],
    67                       loads(dumps(module_test.XYZ_OBJECT)).foo('abc def'))
    68  
    69    def test_nested_class(self):
    70      """Tests that a nested class object is pickled correctly."""
    71      self.assertEqual(
    72          'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum)
    73      self.assertEqual(
    74          'Y:abc',
    75          loads(dumps(module_test.TopClass.MiddleClass.NestedClass('abc'))).datum)
    76  
    77    def test_dynamic_class(self):
    78      """Tests that a nested class object is pickled correctly."""
    79      self.assertEqual(
    80          'Z:abc', loads(dumps(module_test.create_class('abc'))).get())
    81  
    82    def test_generators(self):
    83      with self.assertRaises(TypeError):
    84        dumps((_ for _ in range(10)))
    85  
    86    def test_recursive_class(self):
    87      self.assertEqual(
    88          'RecursiveClass:abc',
    89          loads(dumps(module_test.RecursiveClass('abc').datum)))
    90  
    91    def test_pickle_rlock(self):
    92      rlock_instance = threading.RLock()
    93      rlock_type = type(rlock_instance)
    94  
    95      self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type)
    96  
    97    @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced')
    98    def test_dump_and_load_mapping_proxy(self):
    99      self.assertEqual(
   100          'def', loads(dumps(types.MappingProxyType({'abc': 'def'})))['abc'])
   101      self.assertEqual(
   102          types.MappingProxyType, type(loads(dumps(types.MappingProxyType({})))))
   103  
   104    # pylint: disable=exec-used
   105    @unittest.skipIf(sys.version_info < (3, 7), 'Python 3.7 or above only')
   106    def test_dataclass(self):
   107      exec(
   108          '''
   109  from apache_beam.internal.module_test import DataClass
   110  self.assertEqual(DataClass(datum='abc'), loads(dumps(DataClass(datum='abc'))))
   111      ''')
   112  
   113  
   114  if __name__ == '__main__':
   115    unittest.main()