Handle the non-interactve_env case and import error for unit tests

Author: ian-Liaozy

sdks/python/apache_beam/runners/interactive/display/pipeline_graph.py

@@ -25,10 +25,12 @@
 import collections
 import logging
 import threading
+from typing import Any
 from typing import DefaultDict
 from typing import Dict
 from typing import Iterator
 from typing import List
+from typing import Optional
 from typing import Tuple
 from typing import Union
 
@@ -40,8 +42,12 @@
 
 try:
   import pydot
+  PydotDot = pydot.Dot
+  PydotNode = pydot.Node
+  PydotEdge = pydot.Edge
 except ImportError:
-  pass
+  pydot = None
+  PydotDot = PydotNode = PydotEdge = Any
 
 # pylint does not understand context
 # pylint:disable=dangerous-default-value
@@ -74,7 +80,7 @@ def __init__(
           rendered. See display.pipeline_graph_renderer for available options.
     """
     self._lock = threading.Lock()
-    self._graph: pydot.Dot = None
+    self._graph: Optional[PydotDot] = None  # type: ignore
     self._pipeline_instrument = None
     if isinstance(pipeline, beam.Pipeline):
       self._pipeline_instrument = inst.PipelineInstrument(
@@ -116,7 +122,11 @@ def __init__(
     self._renderer = pipeline_graph_renderer.get_renderer(render_option)
 
   def get_dot(self) -> str:
-    return self._get_graph().to_string()
+    graph: Optional[PydotDot] = self._get_graph()  # type: ignore
+    if graph:
+      return graph.to_string()
+    else:
+      return ""
 
   def display_graph(self):
     """Displays the graph generated."""
@@ -236,6 +246,10 @@ def _construct_graph(
       default_vertex_attrs: (Dict[str, str]) a dict of attributes
       default_edge_attrs: (Dict[str, str]) a dict of attributes
     """
+    if not pydot:
+      self._graph = None
+      return
+
     with self._lock:
       self._graph = pydot.Dot()
 
@@ -270,6 +284,10 @@ def _update_graph(self, vertex_dict=None, edge_dict=None):
           Or (Dict[(str, str), Dict[str, str]]) which maps vertex pairs to edge
           attributes
     """
+    graph: Optional[PydotDot] = self._get_graph()  # type: ignore
+    if not pydot or not graph:
+      return
+
     def set_attrs(ref, attrs):
       for attr_name, attr_val in attrs.items():
         ref.set(attr_name, attr_val)

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

@@ -696,6 +696,9 @@ def test_compute_blocking(self):
     collected = ib.collect(pcoll, raw_records=True)
     self.assertEqual(collected, data)
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   def test_compute_non_blocking(self):
     p = beam.Pipeline(ir.InteractiveRunner())
     data = list(range(5))
@@ -757,6 +760,9 @@ def test_compute_force_recompute(self):
       mock_evict.assert_called_once_with(p)
     self.assertTrue(pcoll in self.env.computed_pcollections)
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   def test_compute_non_blocking_exception(self):
     p = beam.Pipeline(ir.InteractiveRunner())
 
@@ -777,6 +783,9 @@ def raise_error(elem):
     self.assertIsInstance(async_result.exception(), ValueError)
     self.assertFalse(pcoll in self.env.computed_pcollections)
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   @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')
@@ -824,6 +833,9 @@ def test_compute_non_blocking_ipython_widgets(
     mock_clear_output.assert_called_once()
     async_result.result(timeout=60)  # Let it finish
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   def test_compute_dependency_wait_true(self):
     p = beam.Pipeline(ir.InteractiveRunner())
     pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
@@ -855,6 +867,9 @@ def test_compute_dependency_wait_true(self):
       async_res2.result(timeout=60)
       self.assertTrue(pcoll2 in self.env.computed_pcollections)
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   @patch.object(ie.InteractiveEnvironment, 'is_pcollection_computing')
   def test_compute_dependency_wait_false(self, mock_is_computing):
     p = beam.Pipeline(ir.InteractiveRunner())
@@ -878,6 +893,9 @@ def test_compute_dependency_wait_false(self, mock_is_computing):
       spy_execute.assert_called_with({pcoll2}, async_res2, ANY, ANY)
       self.assertTrue(pcoll2 in self.env.computed_pcollections)
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   def test_async_computation_result_cancel(self):
     p = beam.Pipeline(ir.InteractiveRunner())
     # A stream that never finishes to test cancellation
@@ -904,6 +922,9 @@ def test_async_computation_result_cancel(self):
     with self.assertRaises(TimeoutError):
       async_result.result(timeout=1)  # It should not complete successfully
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   @patch(
       'apache_beam.runners.interactive.recording_manager.RecordingManager.'
       '_execute_pipeline_fragment')

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)
@@ -962,6 +965,9 @@ def test_compute_async_blocking(self):
       mock_execute.assert_called_once()
       self.assertTrue(pcoll in ie.current_env().computed_pcollections)
 
+  @unittest.skipIf(
+      not ie.current_env().is_interactive_ready,
+      '[interactive] dependency is not installed.')
   @patch(
       'apache_beam.runners.interactive.recording_manager.AsyncComputationResult'
   )

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

@@ -49,6 +49,18 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+try:
+  from google.cloud.exceptions import ClientError
+  from google.cloud.exceptions import NotFound
+
+  from apache_beam.io.gcp.gcsio import create_storage_client  # pylint: disable=ungrouped-imports
+  GCP_LIBS_AVAILABLE = True
+except ImportError:
+  GCP_LIBS_AVAILABLE = False
+  ClientError = None
+  NotFound = None
+  create_storage_client = None
+
 # Add line breaks to the IPythonLogHandler's HTML output.
 _INTERACTIVE_LOG_STYLE = """
   <style>
@@ -447,11 +459,14 @@ def assert_bucket_exists(bucket_name: str) -> None:
 
     Logs a warning if the bucket cannot be verified to exist.
   """
-  try:
-    from google.cloud.exceptions import ClientError
-    from google.cloud.exceptions import NotFound
+  if not GCP_LIBS_AVAILABLE:
+    _LOGGER.warning(
+        'google-cloud-storage is not installed. '
+        'Unable to verify whether bucket %s exists.',
+        bucket_name)
+    return
 
-    from apache_beam.io.gcp.gcsio import create_storage_client
+  try:
     storage_client = create_storage_client(PipelineOptions())
     storage_client.get_bucket(bucket_name)
   except ClientError as e: