Initial.

Author: claudevdm

sdks/python/apache_beam/internal/pickler.py

@@ -33,9 +33,9 @@
 
 USE_CLOUDPICKLE = 'cloudpickle'
 USE_DILL = 'dill'
-DEFAULT_PICKLE_LIB = USE_DILL
 
-desired_pickle_lib = dill_pickler
+DEFAULT_PICKLE_LIB = USE_CLOUDPICKLE
+desired_pickle_lib = cloudpickle_pickler
 
 
 def dumps(

sdks/python/apache_beam/pipeline_test.py

@@ -287,9 +287,19 @@ def test_no_wait_until_finish(self, mock_info):
   def test_auto_unique_labels(self):
 
     opts = PipelineOptions(["--auto_unique_labels"])
-    with mock.patch.object(uuid, 'uuid4') as mock_uuid_gen:
-      mock_uuids = [mock.Mock(hex='UUID01XXX'), mock.Mock(hex='UUID02XXX')]
-      mock_uuid_gen.side_effect = mock_uuids
+
+    mock_uuids = [mock.Mock(hex='UUID01XXX'), mock.Mock(hex='UUID02XXX')]
+    mock_uuid_gen = mock.Mock(side_effect=mock_uuids)
+
+    original_generate_unique_label = Pipeline._generate_unique_label
+
+    def patched_generate_unique_label(self, transform):
+      with mock.patch.object(uuid, 'uuid4', return_value=mock_uuid_gen()):
+        return original_generate_unique_label(self, transform)
+
+    with mock.patch.object(Pipeline,
+                           '_generate_unique_label',
+                           patched_generate_unique_label):
       with TestPipeline(options=opts) as pipeline:
         pcoll = pipeline | 'pcoll' >> Create([1, 2, 3])
 

sdks/python/apache_beam/runners/direct/direct_runner_test.py

@@ -41,6 +41,15 @@
 from apache_beam.testing import test_pipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
+from apache_beam.utils import shared
+
+
+class DictWrapper(dict):
+  pass
+
+
+def acquire_dict():
+  return DictWrapper()
 
 
 class DirectPipelineResultTest(unittest.TestCase):
@@ -159,18 +168,17 @@ def test_retry_fork_graph(self):
     # currently does not currently support retries.
     p = beam.Pipeline(runner='BundleBasedDirectRunner')
 
-    # TODO(mariagh): Remove the use of globals from the test.
-    global count_b, count_c  # pylint: disable=global-variable-undefined
-    count_b, count_c = 0, 0
+    shared_handler = shared.Shared()
+    counts = shared_handler.acquire(acquire_dict)
+    counts['count_b'] = 0
+    counts['count_c'] = 0
 
     def f_b(x):
-      global count_b  # pylint: disable=global-variable-undefined
-      count_b += 1
+      shared_handler.acquire(acquire_dict)['count_b'] += 1
       raise Exception('exception in f_b')
 
     def f_c(x):
-      global count_c  # pylint: disable=global-variable-undefined
-      count_c += 1
+      shared_handler.acquire(acquire_dict)['count_c'] += 1
       raise Exception('exception in f_c')
 
     names = p | 'CreateNodeA' >> beam.Create(['Ann', 'Joe'])
@@ -180,7 +188,7 @@ def f_c(x):
 
     with self.assertRaises(Exception):
       p.run().wait_until_finish()
-    assert count_b == count_c == 4
+    assert counts['count_b'] == counts['count_c'] == 4
 
   def test_no_partial_writeouts(self):
     class TestTransformEvaluator(_TransformEvaluator):