Fix, add tests

Author: Abacn

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -275,7 +275,7 @@ def load_session(file_path):
     registries = cloudpickle.load(file)
     if type(registries) != dict:
       raise ValueError(
-          "Faled loading session: expected dict, get {}", type(registries))
+          "Faled loading session: expected dict, got {}", type(registries))
     if "coder" in registries:
       typecoders.registry.load_custom_type_coder_tuples(registries["coder"])
     else:

sdks/python/apache_beam/internal/cloudpickle_pickler_test.py

@@ -20,6 +20,7 @@
 # pytype: skip-file
 
 import os
+import tempfile
 import threading
 import types
 import unittest
@@ -31,6 +32,7 @@
 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.typehints.schemas import LogicalTypeRegistry
 from apache_beam.utils import shared
 
 GLOBAL_DICT_REF = module_test.GLOBAL_DICT
@@ -244,6 +246,24 @@ def sample_func():
     unpickled_filename = os.path.abspath(unpickled_code.co_filename)
     self.assertEqual(unpickled_filename, original_filename)
 
+  @mock.patch(
+      "apache_beam.coders.typecoders.registry.load_custom_type_coder_tuples")
+  @mock.patch(
+      "apache_beam.typehints.schemas.LogicalType._known_logical_types.load")
+  def test_dump_load_session(self, logicaltype_mock, coder_mock):
+    session_file = 'pickled'
+
+    with tempfile.TemporaryDirectory() as tmp_dirname:
+      pickled_session_file = os.path.join(tmp_dirname, session_file)
+      beam_cloudpickle.dump_session(pickled_session_file)
+      beam_cloudpickle.load_session(pickled_session_file)
+    load_logical_types = logicaltype_mock.call_args.args
+    load_coders = coder_mock.call_args.args
+    self.assertEqual(len(load_logical_types), 1)
+    self.assertEqual(len(load_coders), 1)
+    self.assertTrue(isinstance(load_logical_types[0], LogicalTypeRegistry))
+    self.assertTrue(isinstance(load_coders[0], list))
+
 
 if __name__ == '__main__':
   unittest.main()

sdks/python/apache_beam/typehints/schemas.py

@@ -720,6 +720,7 @@ def copy_custom(self):
       copy.by_logical_type[logical_type] = urn
       copy.by_language_type[logical_type.language_type()] = logical_type
       copy._custom_urns.add(urn)
+    return copy
 
   def load(self, another):
     self.by_urn.update(another.by_urn)