Fix more linting errors, improve test coverage, make all unit tests pass

Author: ian-Liaozy

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

@@ -24,11 +24,12 @@
 import time
 import unittest
 from typing import NamedTuple
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock, ANY, call
 from concurrent.futures import TimeoutError
 
 import apache_beam as beam
 from apache_beam import dataframe as frames
+from apache_beam.dataframe.frame_base import DeferredBase
 from apache_beam.options.pipeline_options import FlinkRunnerOptions
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.runners.interactive import interactive_beam as ib
@@ -775,6 +776,8 @@ def raise_error(elem):
 
   @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True)
   @patch('apache_beam.runners.interactive.recording_manager.display')
+  @patch('apache_beam.runners.interactive.recording_manager.clear_output')
+  @patch('apache_beam.runners.interactive.recording_manager.HTML')
   @patch('ipywidgets.Button')
   @patch('ipywidgets.FloatProgress')
   @patch('ipywidgets.Output')
@@ -787,6 +790,8 @@ def test_compute_non_blocking_ipython_widgets(
       mock_output,
       mock_progress,
       mock_button,
+      mock_html,
+      mock_clear_output,
       mock_display,
   ):
     self.env._is_in_ipython = True
@@ -795,14 +800,25 @@ def test_compute_non_blocking_ipython_widgets(
     ib.watch(locals())
     self.env.track_user_pipelines()
 
+    mock_controls = mock_vbox.return_value
+    mock_html_instance = mock_html.return_value
+
     async_result = ib.compute(pcoll, blocking=False)
     self.assertIsNotNone(async_result)
     mock_button.assert_called_once_with(description='Cancel')
     mock_progress.assert_called_once()
     mock_output.assert_called_once()
     mock_hbox.assert_called_once()
     mock_vbox.assert_called_once()
-    mock_display.assert_called_once()
+    mock_html.assert_called_once_with('<p>Initializing...</p>')
+
+    self.assertEqual(mock_display.call_count, 2)
+    mock_display.assert_has_calls([
+        call(mock_controls, display_id=async_result._display_id),
+        call(mock_html_instance)
+    ])
+
+    mock_clear_output.assert_called_once()
     async_result.result(timeout=60)  # Let it finish
 
   def test_compute_dependency_wait_true(self):
@@ -885,14 +901,21 @@ def test_async_computation_result_cancel(self):
     with self.assertRaises(TimeoutError):
       async_result.result(timeout=1)  # It should not complete successfully
 
-  def test_compute_multiple_async(self):
+  @patch(
+      'apache_beam.runners.interactive.recording_manager.RecordingManager.'
+      '_execute_pipeline_fragment')
+  def test_compute_multiple_async(self, mock_execute_fragment):
     p = beam.Pipeline(ir.InteractiveRunner())
     pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
     pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6])
     pcoll3 = pcoll1 | 'Map1' >> beam.Map(lambda x: x * 2)
     ib.watch(locals())
     self.env.track_user_pipelines()
 
+    mock_pipeline_result = MagicMock()
+    mock_pipeline_result.state = PipelineState.DONE
+    mock_execute_fragment.return_value = mock_pipeline_result
+
     res1 = ib.compute(pcoll1, blocking=False)
     res2 = ib.compute(pcoll2, blocking=False)
     res3 = ib.compute(pcoll3, blocking=False)  # Depends on pcoll1
@@ -905,9 +928,129 @@ def test_compute_multiple_async(self):
     res2.result(timeout=60)
     res3.result(timeout=60)
 
-    self.assertTrue(pcoll1 in self.env.computed_pcollections)
-    self.assertTrue(pcoll2 in self.env.computed_pcollections)
-    self.assertTrue(pcoll3 in self.env.computed_pcollections)
+    time.sleep(0.1)
+
+    self.assertTrue(
+        pcoll1 in self.env.computed_pcollections, "pcoll1 not marked computed")
+    self.assertTrue(
+        pcoll2 in self.env.computed_pcollections, "pcoll2 not marked computed")
+    self.assertTrue(
+        pcoll3 in self.env.computed_pcollections, "pcoll3 not marked computed")
+
+    self.assertEqual(mock_execute_fragment.call_count, 3)
+
+  @patch(
+      'apache_beam.runners.interactive.interactive_beam.'
+      'deferred_df_to_pcollection')
+  def test_compute_input_flattening(self, mock_deferred_to_pcoll):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    pcoll1 = p | 'C1' >> beam.Create([1])
+    pcoll2 = p | 'C2' >> beam.Create([2])
+    pcoll3 = p | 'C3' >> beam.Create([3])
+    pcoll4 = p | 'C4' >> beam.Create([4])
+
+    class MockDeferred(DeferredBase):
+      def __init__(self, pcoll):
+        mock_expr = MagicMock()
+        super().__init__(mock_expr)
+        self._pcoll = pcoll
+
+      def _get_underlying_pcollection(self):
+        return self._pcoll
+
+    deferred_pcoll = MockDeferred(pcoll4)
+
+    mock_deferred_to_pcoll.return_value = (pcoll4, p)
+
+    ib.watch(locals())
+    self.env.track_user_pipelines()
+
+    with patch.object(self.env, 'get_recording_manager') as mock_get_rm:
+      mock_rm = MagicMock()
+      mock_get_rm.return_value = mock_rm
+      ib.compute(pcoll1, [pcoll2], {'a': pcoll3}, deferred_pcoll)
+
+      expected_pcolls = {pcoll1, pcoll2, pcoll3, pcoll4}
+      mock_rm.compute_async.assert_called_once_with(
+          expected_pcolls,
+          wait_for_inputs=True,
+          blocking=False,
+          runner=None,
+          options=None,
+          force_compute=False)
+
+  def test_compute_invalid_input_type(self):
+    with self.assertRaisesRegex(ValueError,
+                                "not a dict, an iterable or a PCollection"):
+      ib.compute(123)
+
+  def test_compute_mixed_pipelines(self):
+    p1 = beam.Pipeline(ir.InteractiveRunner())
+    pcoll1 = p1 | 'C1' >> beam.Create([1])
+    p2 = beam.Pipeline(ir.InteractiveRunner())
+    pcoll2 = p2 | 'C2' >> beam.Create([2])
+    ib.watch(locals())
+    self.env.track_user_pipelines()
+
+    with self.assertRaisesRegex(
+        ValueError, "All PCollections must belong to the same pipeline"):
+      ib.compute(pcoll1, pcoll2)
+
+  @patch(
+      'apache_beam.runners.interactive.interactive_beam.'
+      'deferred_df_to_pcollection')
+  @patch.object(ib, 'watch')
+  def test_compute_with_deferred_base(self, mock_watch, mock_deferred_to_pcoll):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    pcoll = p | 'C1' >> beam.Create([1])
+
+    class MockDeferred(DeferredBase):
+      def __init__(self, pcoll):
+        # Provide a dummy expression to satisfy DeferredBase.__init__
+        mock_expr = MagicMock()
+        super().__init__(mock_expr)
+        self._pcoll = pcoll
+
+      def _get_underlying_pcollection(self):
+        return self._pcoll
+
+    deferred = MockDeferred(pcoll)
+
+    mock_deferred_to_pcoll.return_value = (pcoll, p)
+
+    with patch.object(self.env, 'get_recording_manager') as mock_get_rm:
+      mock_rm = MagicMock()
+      mock_get_rm.return_value = mock_rm
+      ib.compute(deferred)
+
+      mock_deferred_to_pcoll.assert_called_once_with(deferred)
+      self.assertEqual(mock_watch.call_count, 2)
+      mock_watch.assert_has_calls([
+          call({f'anonymous_pcollection_{id(pcoll)}': pcoll}),
+          call({f'anonymous_pipeline_{id(p)}': p})
+      ],
+                                  any_order=False)
+      mock_rm.compute_async.assert_called_once_with({pcoll},
+                                                    wait_for_inputs=True,
+                                                    blocking=False,
+                                                    runner=None,
+                                                    options=None,
+                                                    force_compute=False)
+
+  def test_compute_new_pipeline(self):
+    p = beam.Pipeline(ir.InteractiveRunner())
+    pcoll = p | 'Create' >> beam.Create([1])
+    # NOT calling ib.watch() or track_user_pipelines()
+
+    with patch.object(self.env, 'get_recording_manager') as mock_get_rm, \
+        patch.object(ib, 'watch') as mock_watch:
+      mock_rm = MagicMock()
+      mock_get_rm.return_value = mock_rm
+      ib.compute(pcoll)
+
+      mock_watch.assert_called_with({f'anonymous_pipeline_{id(p)}': p})
+      mock_get_rm.assert_called_once_with(p, create_if_absent=True)
+      mock_rm.compute_async.assert_called_once()
 
 
 if __name__ == '__main__':

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

@@ -163,16 +163,16 @@ def _on_done(self, future: Future):
     exc = future.exception()
     if exc:
       self.update_display(f'Error: {exc}', 1.0)
-      _LOGGER.error(f'Asynchronous computation failed: {exc}', exc_info=exc)
+      _LOGGER.error('Asynchronous computation failed: %s', exc, exc_info=exc)
     else:
       self.update_display('Computation Finished Successfully.', 1.0)
       res = future.result()
       if res and res.state == PipelineState.DONE:
         self._env.mark_pcollection_computed(self._pcolls)
       else:
         _LOGGER.warning(
-            'Async computation finished but state is not DONE:'
-            f" {res.state if res else 'Unknown'}")
+            'Async computation finished but state is not DONE: %s',
+            res.state if res else 'Unknown')
 
   def cancel(self):
     if self._future.done():
@@ -198,8 +198,8 @@ def cancel(self):
         # The future will be cancelled by the runner if successful
         return True
       except Exception as e:
-        self.update_display(f'Error sending cancel signal: {e}')
-        _LOGGER.warning(f'Error during pipeline cancel(): {e}', exc_info=e)
+        self.update_display('Error sending cancel signal: %s', e)
+        _LOGGER.warning('Error during pipeline cancel(): %s', e, exc_info=e)
         # Still try to cancel the future as a fallback
         return self._future.cancel()
     else:
@@ -495,8 +495,8 @@ def _run_async_computation(
 
       self._env.mark_pcollection_computing(pcolls_to_compute)
       _LOGGER.info(
-          'Starting asynchronous computation for'
-          f' {len(pcolls_to_compute)} PCollections.')
+          'Starting asynchronous computation for %d PCollections.',
+          len(pcolls_to_compute))
 
       pipeline_result = self._execute_pipeline_fragment(
           pcolls_to_compute, async_result, runner, options)
@@ -515,7 +515,7 @@ def _run_async_computation(
       #   )
       return pipeline_result
     except Exception as e:
-      _LOGGER.exception(f'Exception during asynchronous computation: {e}')
+      _LOGGER.exception('Exception during asynchronous computation: %s', e)
       raise
     # finally:
     #   self._env.unmark_pcollection_computing(pcolls_to_compute)
@@ -685,9 +685,9 @@ def compute_async(
           self._env.mark_pcollection_computed(pcolls_to_compute)
         else:
           _LOGGER.error(
-              f'Blocking computation failed. State: {pipeline_result.state}')
+              'Blocking computation failed. State: %s', pipeline_result.state)
           raise RuntimeError(
-              f'Blocking computation failed. State: {pipeline_result.state}')
+              'Blocking computation failed. State: %s', pipeline_result.state)
       finally:
         self._env.unmark_pcollection_computing(pcolls_to_compute)
       return None
@@ -715,15 +715,13 @@ def _get_pcoll_id_map(self):
     pcoll_to_id = {}
     if self._pipeline_graph._pipeline_instrument:
       pcoll_to_id = self._pipeline_graph._pipeline_instrument._pcoll_to_pcoll_id
-    else:
-      # Fallback for proto-based PipelineGraph, though less likely in this context
-      proto = self._pipeline_graph._pipeline_proto
     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."""
+    """Gets all upstream PCollection dependencies
+    for the given set of PCollections."""
     if not self._pipeline_graph:
       return set()
 
@@ -732,7 +730,6 @@ def _get_all_dependencies(
       return set()
 
     pcoll_to_id = analyzer._pcoll_to_pcoll_id
-    id_to_pcoll_str = {v: k for k, v in pcoll_to_id.items()}
 
     target_pcoll_ids = {
         pcoll_to_id.get(str(pcoll))
@@ -784,7 +781,8 @@ def _wait_for_dependencies(
       pcolls: Set[beam.pvalue.PCollection],
       async_result: Optional[AsyncComputationResult] = None,
   ) -> bool:
-    """Waits for any dependencies of the given PCollections that are currently being computed."""
+    """Waits for any dependencies of the given
+    PCollections that are currently being computed."""
     dependencies = self._get_all_dependencies(pcolls)
     computing_deps: Dict[beam.pvalue.PCollection, AsyncComputationResult] = {}
 
@@ -800,10 +798,11 @@ def _wait_for_dependencies(
 
     if async_result:
       async_result.update_display(
-          f'Waiting for {len(computing_deps)} dependencies to finish...')
+          'Waiting for %d dependencies to finish...', len(computing_deps))
     _LOGGER.info(
-        f'Waiting for {len(computing_deps)} dependencies:'
-        f' {computing_deps.keys()}')
+        'Waiting for %d dependencies: %s',
+        len(computing_deps),
+        computing_deps.keys())
 
     futures_to_wait = list(
         set(comp._future for comp in computing_deps.values()))
@@ -823,7 +822,7 @@ def _wait_for_dependencies(
     except Exception as e:
       if async_result:
         async_result.update_display(f'Dependency failed: {e}')
-      _LOGGER.error(f'Dependency computation failed: {e}', exc_info=e)
+      _LOGGER.error('Dependency computation failed: %s', e, exc_info=e)
       return False
 
   def record(