Add tests.

Author: claudevdm

sdks/python/apache_beam/internal/cloudpickle_pickler_test.py

@@ -27,12 +27,83 @@
 from apache_beam.internal import module_test
 from apache_beam.internal.cloudpickle_pickler import dumps
 from apache_beam.internal.cloudpickle_pickler import loads
+from apache_beam.utils import shared
+
+GLOBAL_DICT_REF = module_test.GLOBAL_DICT
+
+
+# Allow weakref to dict
+class DictWrapper(dict):
+  pass
+
+
+MAIN_MODULE_DICT = DictWrapper()
+
+
+def acquire_dict():
+  return DictWrapper()
 
 
 class PicklerTest(unittest.TestCase):
 
   NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
 
+  def test_module_global_pickled_by_value(self):
+    self.assertIsNot(
+        module_test.GLOBAL_DICT, loads(dumps(module_test.GLOBAL_DICT)))
+
+  def test_main_shared_global_pickled_by_ref(self):
+    shared_handler = shared.Shared()
+    original_dict = shared_handler.acquire(acquire_dict)
+
+    unpickled_dict = loads(
+        dumps(lambda: shared_handler.acquire(acquire_dict)))()
+
+    self.assertIs(original_dict, unpickled_dict)
+
+  def test_function_defined_in_main_pickles_module_global_by_ref(self):
+    def returns_global_dict():
+      return module_test.GLOBAL_DICT
+
+    self.assertIs(module_test.GLOBAL_DICT, loads(dumps(returns_global_dict))())
+
+  def test_function_defined_in_main_pickles_module_global_ref_by_value(self):
+    def returns_global_dict():
+      return GLOBAL_DICT_REF
+
+    self.assertIsNot(
+        module_test.GLOBAL_DICT, loads(dumps(returns_global_dict))())
+
+  def test_fn_contains_unpicklable_by_ref(self):
+    self.assertEqual(
+        module_test.UNPICLABLE_INSTANCE,
+        loads(dumps(module_test.fn_returns_unpicklable))())
+
+  def test_imported_closure_contains_unpicklable_by_value_fails(self):
+    # The entire closure is pickled by value, and therefore module_test is
+    # not imported. Requires the module global to be pickled by value.
+    with self.assertRaises(Exception):
+      loads(dumps(module_test.closure_contains_unpicklable()))
+
+  def test_imported_closure_contains_unpicklable_imports_self(self):
+    # The closure imports module_test within the function definition
+    # and returns self.UNPICLABLE_INSTANCE. This allows cloudpickle
+    # to use submimort to reference module_test.UNPICLABLE_INSTANCE
+    self.assertIs(
+        module_test.UNPICLABLE_INSTANCE,
+        loads(dumps(module_test.closure_contains_unpicklable_imports_self()))())
+
+  def test_main_closure_contains_imported_unpicklable_by_ref(self):
+    def outer():
+      def inner():
+        return module_test.UNPICLABLE_INSTANCE
+
+      return inner
+
+    # Uses subimport to reference module_test.UNPICLABLE_INSTANCE rather than
+    # recreate.
+    loads(dumps(outer()))
+
   def test_pickle_nested_enum_descriptor(self):
     NestedEnum = proto2_coder_test_messages_pb2.MessageD.NestedEnum
 

sdks/python/apache_beam/internal/module_test.py

@@ -23,6 +23,38 @@
 import sys
 from typing import Any
 
+GLOBAL_DICT = {}
+
+
+class UnPicklable:
+  def __init__(self, x):
+    self.x = x
+
+  def __reduce__(self):
+    return NotImplemented
+
+
+UNPICLABLE_INSTANCE = UnPicklable(5)
+
+
+def closure_contains_unpicklable():
+  def inner():
+    return UNPICLABLE_INSTANCE
+
+  return inner
+
+
+def closure_contains_unpicklable_imports_self():
+  def inner():
+    import apache_beam.internal.module_test as self_module
+    return self_module.UNPICLABLE_INSTANCE
+
+  return inner
+
+
+def fn_returns_unpicklable():
+  return UNPICLABLE_INSTANCE
+
 
 class TopClass(object):
   class NestedClass(object):