diff --git a/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py b/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py
index 8eefec744879..98deb7f34636 100644
--- a/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py
+++ b/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,6 +49,7 @@
 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
@@ -54,9 +57,11 @@ def setUp(self):
     # 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()
 
diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini
index d22751996aec..558e9a743b7c 100644
--- a/sdks/python/tox.ini
+++ b/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}"