Fix log fixtures

Signed-off-by: pdmurray <peynmurray@gmail.com>

python/ray/tests/test_widgets.py

@@ -2,9 +2,17 @@
 from unittest import mock
 
 import pytest
+import ray
 from ray.widgets.util import repr_with_fallback, _can_display_ipywidgets
 
 
+@pytest.fixture
+def logs(propagate_logs, caplog):
+    """A log fixture which captures logs during a test."""
+    caplog.set_level(logging.INFO)
+    return caplog
+
+
 @pytest.fixture
 def fancy_mimebundle():
     return {
@@ -13,26 +21,24 @@ def fancy_mimebundle():
     }
 
 
-@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch.object(ray.widgets.util, "in_notebook")
 @mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
 def test_repr_with_fallback_missing(
     mock_import_module,
     mock_find_spec,
-    mock_get_ipython_shell_name,
-    propagate_logs,
-    caplog,
+    mock_in_notebook,
+    logs,
     fancy_mimebundle,
 ):
     """Test that missing notebook dependencies trigger a log message."""
-    caplog.set_level(logging.INFO)
 
     def raise_import_error(*args):
         raise ImportError
 
     mock_import_module.side_effect = raise_import_error
     mock_find_spec.return_value = None
-    mock_get_ipython_shell_name.return_value = "ZMQInteractiveShell"
+    mock_in_notebook.return_value = True
 
     class DummyObject:
         def __repr__(self):
@@ -44,29 +50,27 @@ def _repr_mimebundle_(self, **kwargs):
 
     result = DummyObject()._repr_mimebundle_()
     assert result == {"text/plain": "dummy repr"}
-    assert "Missing packages:" in caplog.records[-1].msg
+    assert "Missing packages:" in logs.records[-1].msg
 
 
-@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch.object(ray.widgets.util, "in_notebook")
 @mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
 def test_repr_with_fallback_outdated(
     mock_import_module,
     mock_find_spec,
-    mock_get_ipython_shell_name,
-    propagate_logs,
-    caplog,
+    mock_in_notebook,
+    logs,
     fancy_mimebundle,
 ):
     """Test that outdated notebook dependencies trigger a log message."""
-    caplog.set_level(logging.INFO)
 
     class MockDep:
         __version__ = "7.0.0"
 
     mock_import_module.return_value = MockDep()
     mock_find_spec.return_value = "a valid import spec"
-    mock_get_ipython_shell_name.return_value = "ZMQInteractiveShell"
+    mock_in_notebook.return_value = True
 
     class DummyObject:
         def __repr__(self):
@@ -78,29 +82,27 @@ def _repr_mimebundle_(self, **kwargs):
 
     result = DummyObject()._repr_mimebundle_()
     assert result == {"text/plain": "dummy repr"}
-    assert "Outdated packages:" in caplog.records[-1].msg
+    assert "Outdated packages:" in logs.records[-1].msg
 
 
-@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch.object(ray.widgets.util, "_can_display_ipywidgets")
 @mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
 def test_repr_with_fallback_valid(
     mock_import_module,
     mock_find_spec,
-    mock_get_ipython_shell_name,
-    propagate_logs,
-    caplog,
+    mock_can_display_ipywidgets,
+    logs,
     fancy_mimebundle,
 ):
     """Test that valid notebook dependencies don't trigger a log message."""
-    caplog.set_level(logging.INFO)
 
     class MockDep:
         __version__ = "8.0.0"
 
     mock_import_module.return_value = MockDep()
     mock_find_spec.return_value = "a valid import spec"
-    mock_get_ipython_shell_name.return_value = "ZMQInteractiveShell"
+    mock_can_display_ipywidgets.return_value = True
 
     class DummyObject:
         def __repr__(self):
@@ -111,22 +113,21 @@ def _repr_mimebundle_(self, **kwargs):
             return fancy_mimebundle
 
     result = DummyObject()._repr_mimebundle_()
-    assert len(caplog.records) == 0
+    assert len(logs.records) == 0
     assert result == fancy_mimebundle
 
 
-@mock.patch("ray.widgets.util._can_display_ipywidgets")
+@mock.patch.object(ray.widgets.util, "_can_display_ipywidgets")
 @mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
 def test_repr_with_fallback_invalid_shell(
     mock_import_module,
     mock_find_spec,
     mock_can_display_ipywidgets,
-    caplog,
+    logs,
     fancy_mimebundle,
 ):
     """Test that the mimebundle is correctly stripped if run in an invalid shell."""
-    caplog.set_level(logging.INFO)
 
     class MockDep:
         __version__ = "8.0.0"
@@ -144,11 +145,11 @@ def _repr_mimebundle_(self, **kwargs):
             return fancy_mimebundle
 
     result = DummyObject()._repr_mimebundle_()
-    assert len(caplog.records) == 0
+    assert len(logs.records) == 0
     assert result == {"text/plain": "dummy repr"}
 
 
-@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch.object(ray.widgets.util, "_get_ipython_shell_name")
 @mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
 @pytest.mark.parametrize(
@@ -175,6 +176,7 @@ class MockDep:
     mock_get_ipython_shell_name.return_value = shell
 
     assert _can_display_ipywidgets(["somedep", "8"], message="") == can_display
+    mock_get_ipython_shell_name.assert_called()
 
 
 if __name__ == "__main__":

python/ray/widgets/util.py

@@ -179,9 +179,9 @@ def wrapped(self, *args, **kwargs):
 
 def _get_ipython_shell_name() -> str:
     if "IPython" in sys.modules:
-        import IPython
+        from IPython import get_ipython
 
-        return IPython.get_ipython().__class__.__name__
+        return get_ipython().__class__.__name__
     return ""