[BEAM-8837] Fix pcoll_visualization tests

1. Added [interactive] as extras for py36 and py37 gcp tests so they
will be executed at least once during the pre-commit. The [interactive]
dependencies were removed after PR apache#10227 deprecating the old setup config,
thus not triggering any pcoll_visualization tests since then.
2. Fixed the Setup routine of pcoll_visualization_tests so that global
environment is always fresh no matter how the unit tests modify its
states. The bug was caused by wrong execution paths instructed by
undeterministic tests execution sequence. Somehow, the unittest's patch
utility fails to patch _to_element_list some times in multi-threading.
Using InteractiveRunner when building pipelines to lazily initialize a
cache manager for those scenarios.

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

@@ -24,7 +24,9 @@
 
 import apache_beam as beam
 from apache_beam.runners import runner
+from apache_beam.runners.interactive import interactive_beam as ib
 from apache_beam.runners.interactive import interactive_environment as ie
+from apache_beam.runners.interactive import interactive_runner as ir
 from apache_beam.runners.interactive.display import pcoll_visualization as pv
 
 # TODO(BEAM-8288): clean up the work-around of nose tests using Python2 without
@@ -47,16 +49,19 @@
 class PCollectionVisualizationTest(unittest.TestCase):
 
   def setUp(self):
+    ie.new_env()
     # Allow unit test to run outside of ipython kernel since we don't test the
     # frontend rendering in unit tests.
     pv._pcoll_visualization_ready = True
     # Generally test the logic where notebook is connected to the assumed
     # ipython kernel by forcefully setting notebook check to True.
     ie.current_env()._is_in_notebook = True
 
-    self._p = beam.Pipeline()
+    self._p = beam.Pipeline(ir.InteractiveRunner())
+    ib.watch(self)
     # pylint: disable=range-builtin-not-iterating
     self._pcoll = self._p | 'Create' >> beam.Create(range(1000))
+    ib.watch(self)
 
   def test_raise_error_for_non_pcoll_input(self):
     class Foo(object):
@@ -74,19 +79,14 @@ def test_pcoll_visualization_generate_unique_display_id(self):
     self.assertNotEqual(pv_1._overview_display_id, pv_2._overview_display_id)
     self.assertNotEqual(pv_1._df_display_id, pv_2._df_display_id)
 
+  def _mock_to_element_list(self):
+    return [1, 2, 3, 4, 5, 6, 7, 8]
+
   @patch('apache_beam.runners.interactive.display.pcoll_visualization'
-         '.PCollectionVisualization._to_element_list', lambda x: [1, 2, 3])
+         '.PCollectionVisualization._to_element_list', _mock_to_element_list)
   def test_one_shot_visualization_not_return_handle(self):
     self.assertIsNone(pv.visualize(self._pcoll))
 
-  def _mock_to_element_list(self):
-    yield [1, 2, 3]
-    yield [1, 2, 3, 4]
-    yield [1, 2, 3, 4, 5]
-    yield [1, 2, 3, 4, 5, 6]
-    yield [1, 2, 3, 4, 5, 6, 7]
-    yield [1, 2, 3, 4, 5, 6, 7, 8]
-
   @patch('apache_beam.runners.interactive.display.pcoll_visualization'
          '.PCollectionVisualization._to_element_list', _mock_to_element_list)
   def test_dynamic_plotting_return_handle(self):
@@ -139,11 +139,12 @@ def test_auto_stop_dynamic_plotting_when_job_is_terminated(
     mocked_timeloop.assert_called()
 
   @patch('apache_beam.runners.interactive.display.pcoll_visualization'
-         '.PCollectionVisualization._to_element_list', lambda x: [1, 2, 3])
+         '.PCollectionVisualization._to_element_list', _mock_to_element_list)
   @patch('pandas.DataFrame.sample')
   def test_display_plain_text_when_kernel_has_no_frontend(self,
                                                           _mocked_sample):
-    ie.new_env()  # Resets the notebook check. Should be False in unit tests.
+    # Resets the notebook check to False.
+    ie.current_env()._is_in_notebook = False
     self.assertIsNone(pv.visualize(self._pcoll))
     _mocked_sample.assert_called_once()
 

sdks/python/tox.ini

@@ -231,28 +231,28 @@ commands =
 [testenv:py36-gcp]
 setenv =
   RUN_SKIPPED_PY3_TESTS=0
-extras = test,gcp
+extras = test,gcp,interactive
 commands =
   python setup.py nosetests --ignore-files '.*py3[7-9]\.py$' {posargs}
 
 [testenv:py36-gcp-pytest]
 setenv =
   RUN_SKIPPED_PY3_TESTS=0
-extras = test,gcp
+extras = test,gcp,interactive
 commands =
   {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
 
 [testenv:py37-gcp]
 setenv =
   RUN_SKIPPED_PY3_TESTS=0
-extras = test,gcp
+extras = test,gcp,interactive
 commands =
   python setup.py nosetests --ignore-files '.*py3[8-9]\.py$' {posargs}
 
 [testenv:py37-gcp-pytest]
 setenv =
   RUN_SKIPPED_PY3_TESTS=0
-extras = test,gcp
+extras = test,gcp,interactive
 commands =
   {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"