Skip some tests for non-interactve_env and fix errors in unit tests

Author: ian-Liaozy

sdks/python/apache_beam/runners/interactive/interactive_beam_test.py

@@ -71,6 +71,9 @@ def _get_watched_pcollections_with_variable_names():
   return watched_pcollections
 
 
+@unittest.skipIf(
+    not ie.current_env().is_interactive_ready,
+    '[interactive] dependency is not installed.')
 @isolated_env
 class InteractiveBeamTest(unittest.TestCase):
   def setUp(self):
@@ -677,6 +680,9 @@ def test_default_value_for_invalid_worker_number(self):
     self.assertEqual(meta.num_workers, 2)
 
 
+@unittest.skipIf(
+    not ie.current_env().is_interactive_ready,
+    '[interactive] dependency is not installed.')
 @isolated_env
 class InteractiveBeamComputeTest(unittest.TestCase):
   def setUp(self):

sdks/python/apache_beam/runners/interactive/interactive_environment_test.py

@@ -34,6 +34,9 @@
 _module_name = 'apache_beam.runners.interactive.interactive_environment_test'
 
 
+@unittest.skipIf(
+    not ie.current_env().is_interactive_ready,
+    '[interactive] dependency is not installed.')
 @isolated_env
 class InteractiveEnvironmentTest(unittest.TestCase):
   def setUp(self):

sdks/python/apache_beam/runners/interactive/recording_manager.py

@@ -438,7 +438,7 @@ def __init__(
     self._executor = ThreadPoolExecutor(max_workers=os.cpu_count())
     self._env = ie.current_env()
     self._async_computations: dict[str, AsyncComputationResult] = {}
-    self._pipeline_graph = PipelineGraph(self.user_pipeline)
+    self._pipeline_graph = None
 
   def _execute_pipeline_fragment(
       self,
@@ -492,7 +492,6 @@ def _run_async_computation(
         if not self._wait_for_dependencies(pcolls_to_compute, async_result):
           raise RuntimeError('Dependency computation failed or was cancelled.')
 
-      self._env.mark_pcollection_computing(pcolls_to_compute)
       _LOGGER.info(
           'Starting asynchronous computation for %d PCollections.',
           len(pcolls_to_compute))
@@ -696,6 +695,7 @@ def compute_async(
       async_result = AsyncComputationResult(
           future, pcolls_to_compute, self.user_pipeline, self)
       self._async_computations[async_result._display_id] = async_result
+      self._env.mark_pcollection_computing(pcolls_to_compute)
 
       def task():
         try:
@@ -709,22 +709,39 @@ def task():
       self._executor.submit(task)
       return async_result
 
+  def _get_pipeline_graph(self):
+    """Lazily initializes and returns the PipelineGraph."""
+    if self._pipeline_graph is None:
+      try:
+        # Try to create the graph.
+        self._pipeline_graph = PipelineGraph(self.user_pipeline)
+      except (ImportError, NameError, AttributeError):
+        # If pydot is missing, PipelineGraph() might crash.
+        _LOGGER.warning(
+            "Could not create PipelineGraph (pydot missing?). " \
+            "Async features disabled."
+        )
+        self._pipeline_graph = None
+    return self._pipeline_graph
+
   def _get_pcoll_id_map(self):
     """Creates a map from PCollection object to its ID in the proto."""
     pcoll_to_id = {}
-    if self._pipeline_graph._pipeline_instrument:
-      pcoll_to_id = self._pipeline_graph._pipeline_instrument._pcoll_to_pcoll_id
+    graph = self._get_pipeline_graph()
+    if graph and graph._pipeline_instrument:
+      pcoll_to_id = graph._pipeline_instrument._pcoll_to_pcoll_id
     return {v: k for k, v in pcoll_to_id.items()}
 
   def _get_all_dependencies(
       self,
       pcolls: set[beam.pvalue.PCollection]) -> set[beam.pvalue.PCollection]:
     """Gets all upstream PCollection dependencies
     for the given set of PCollections."""
-    if not self._pipeline_graph:
+    graph = self._get_pipeline_graph()
+    if not graph:
       return set()
 
-    analyzer = self._pipeline_graph._pipeline_instrument
+    analyzer = graph._pipeline_instrument
     if not analyzer:
       return set()
 
@@ -751,8 +768,8 @@ def _get_all_dependencies(
     queue = collections.deque(target_pcoll_ids)
     visited_pcoll_ids = set(target_pcoll_ids)
 
-    producers = self._pipeline_graph._producers
-    transforms = self._pipeline_graph._pipeline_proto.components.transforms
+    producers = graph._producers
+    transforms = graph._pipeline_proto.components.transforms
 
     while queue:
       pcoll_id = queue.popleft()

sdks/python/apache_beam/runners/interactive/recording_manager_test.py

@@ -47,6 +47,9 @@
 from apache_beam.utils.windowed_value import WindowedValue
 
 
+@unittest.skipIf(
+    not ie.current_env().is_interactive_ready,
+    '[interactive] dependency is not installed.')
 class AsyncComputationResultTest(unittest.TestCase):
   def setUp(self):
     self.mock_future = MagicMock(spec=Future)
@@ -664,6 +667,9 @@ def test_describe(self):
         cache_manager.size('full', letters_stream.cache_key))
 
 
+@unittest.skipIf(
+    not ie.current_env().is_interactive_ready,
+    '[interactive] dependency is not installed.')
 class RecordingManagerTest(unittest.TestCase):
   def test_basic_execution(self):
     """A basic pipeline to be used as a smoke test."""
@@ -987,12 +993,6 @@ def capture_task(task):
 
     mock_submit.side_effect = capture_task
 
-    res = rm.compute_async({pcoll}, blocking=False)
-    self.assertIs(res, mock_async_res_instance)
-    mock_submit.assert_called_once()
-    self.assertIsNotNone(task_submitted)
-
-    # Patch dependencies of _run_async_computation
     with patch.object(
         rm, '_wait_for_dependencies', return_value=True
     ), patch.object(
@@ -1002,10 +1002,22 @@ def capture_task(task):
         'mark_pcollection_computing',
         wraps=ie.current_env().mark_pcollection_computing,
     ) as wrapped_mark:
-      # Run the task to trigger the marks
-      task_submitted()
+
+      res = rm.compute_async({pcoll}, blocking=False)
       wrapped_mark.assert_called_once_with({pcoll})
 
+    # Run the task to trigger the marks
+    self.assertIs(res, mock_async_res_instance)
+    mock_submit.assert_called_once()
+    self.assertIsNotNone(task_submitted)
+
+    with patch.object(
+        rm, '_wait_for_dependencies', return_value=True
+    ), patch.object(
+        rm, '_execute_pipeline_fragment'
+    ) as _:
+      task_submitted()
+
     self.assertTrue(pcoll in ie.current_env().computing_pcollections)
 
   def test_get_all_dependencies(self):

sdks/python/apache_beam/runners/interactive/utils_test.py

@@ -244,6 +244,9 @@ def test_child_module_logger_can_override_logging_level(self, mock_emit):
     reason='[interactive] dependency is not installed.')
 class ProgressIndicatorTest(unittest.TestCase):
   def setUp(self):
+    self.gcs_patcher = patch(
+        'apache_beam.io.gcp.gcsfilesystem.GCSFileSystem.delete')
+    self.gcs_patcher.start()
     ie.new_env()
 
   @patch('IPython.get_ipython', new_callable=mock_get_ipython)
@@ -279,6 +282,9 @@ def test_progress_in_HTML_JS_when_in_notebook(
         mocked_html.assert_called()
       mocked_js.assert_called()
 
+  def tearDown(self):
+    self.gcs_patcher.stop()
+
 
 @unittest.skipIf(
     not ie.current_env().is_interactive_ready,
@@ -287,6 +293,9 @@ class MessagingUtilTest(unittest.TestCase):
   SAMPLE_DATA = {'a': [1, 2, 3], 'b': 4, 'c': '5', 'd': {'e': 'f'}}
 
   def setUp(self):
+    self.gcs_patcher = patch(
+        'apache_beam.io.gcp.gcsfilesystem.GCSFileSystem.delete')
+    self.gcs_patcher.start()
     ie.new_env()
 
   def test_as_json_decorator(self):
@@ -298,6 +307,9 @@ def dummy():
     # dictionaries remember the order of items inserted.
     self.assertEqual(json.loads(dummy()), MessagingUtilTest.SAMPLE_DATA)
 
+  def tearDown(self):
+    self.gcs_patcher.stop()
+
 
 class GeneralUtilTest(unittest.TestCase):
   def test_pcoll_by_name(self):