test: catch expected UserWarning in test_repr_mimebundle_selection_logic (#17717)

Assert expected `UserWarning` during the fallback path test in
`test_repr_mimebundle_selection_logic`.
When testing fallback behavior on `anywidget` failure,
`html.repr_mimebundle` emits a warning (`Anywidget mode is not available
or failed to load`). This change uses `pytest.warns(UserWarning,
match=...)` to handle the warning cleanly in unit tests.

Fixes #<534917716> 🦕

Author: shuoweil

packages/bigframes/tests/unit/display/test_anywidget.py

@@ -45,11 +45,13 @@ def test_navigation_to_invalid_page_resets_to_valid_page_without_deadlock():
     mock_df._block = mock_block
 
     # We mock _initial_load to avoid complex setup
-    with mock.patch.object(TableWidget, "_initial_load"):
-        with bigframes.option_context(
+    with (
+        mock.patch.object(TableWidget, "_initial_load"),
+        bigframes.option_context(
             "display.render_mode", "anywidget", "display.max_rows", 10
-        ):
-            widget = TableWidget(mock_df)
+        ),
+    ):
+        widget = TableWidget(mock_df)
 
     # Simulate "loaded data but unknown total rows" state
     widget.page_size = 10
@@ -179,11 +181,11 @@ def test_page_size_change_resets_sort(mock_df):
 
 def test_cell_execution_count_propagation(mock_df):
     """Test that the captured cell_execution_count is propagated to to_pandas_batches."""
-    with mock.patch(
-        "bigframes.core.utils.get_ipython_execution_count", return_value=42
+    with (
+        mock.patch("bigframes.core.utils.get_ipython_execution_count", return_value=42),
+        bigframes.option_context("display.render_mode", "anywidget"),
     ):
-        with bigframes.option_context("display.render_mode", "anywidget"):
-            widget = TableWidget(mock_df)
+        widget = TableWidget(mock_df)
 
     assert widget._cell_execution_count == 42
 
@@ -274,19 +276,23 @@ def __init__(self):
 
 
 def test_init_raises_if_anywidget_not_installed():
-    with mock.patch("bigframes.display.anywidget._ANYWIDGET_INSTALLED", False):
-        with pytest.raises(ImportError):
-            from bigframes.display.anywidget import TableWidget
+    with (
+        mock.patch("bigframes.display.anywidget._ANYWIDGET_INSTALLED", False),
+        pytest.raises(ImportError),
+    ):
+        from bigframes.display.anywidget import TableWidget
 
-            TableWidget(mock.Mock())
+        TableWidget(mock.Mock())
 
 
 def test_init_initializes_attributes(mock_df_deferred):
     from bigframes.display.anywidget import TableWidget
 
-    with bigframes.option_context("display.render_mode", "anywidget"):
-        with mock.patch.object(TableWidget, "_initial_load"):
-            widget = TableWidget(mock_df_deferred)
+    with (
+        bigframes.option_context("display.render_mode", "anywidget"),
+        mock.patch.object(TableWidget, "_initial_load"),
+    ):
+        widget = TableWidget(mock_df_deferred)
 
     assert widget._dataframe is mock_df_deferred
     assert widget.page == 0
@@ -328,19 +334,21 @@ def test_validate_page_clamping(mock_df_deferred):
 def test_validate_page_size(mock_df_deferred):
     from bigframes.display.anywidget import TableWidget
 
-    with bigframes.option_context("display.render_mode", "anywidget"):
-        with mock.patch.object(TableWidget, "_initial_load"):
-            widget = TableWidget(mock_df_deferred)
+    with (
+        bigframes.option_context("display.render_mode", "anywidget"),
+        mock.patch.object(TableWidget, "_initial_load"),
+    ):
+        widget = TableWidget(mock_df_deferred)
 
-            widget.page_size = 50
-            assert widget.page_size == 50
+        widget.page_size = 50
+        assert widget.page_size == 50
 
-            original_size = widget.page_size
-            widget.page_size = -5
-            assert widget.page_size == original_size
+        original_size = widget.page_size
+        widget.page_size = -5
+        assert widget.page_size == original_size
 
-            widget.page_size = 10000
-            assert widget.page_size == 1000
+        widget.page_size = 10000
+        assert widget.page_size == 1000
 
 
 def test_page_size_change_resets_page_and_sort(mock_df_deferred):
@@ -373,14 +381,16 @@ def test_page_size_change_resets_batches(mock_df_deferred):
 def test_sort_change_resets_batches(mock_df_deferred):
     from bigframes.display.anywidget import TableWidget
 
-    with bigframes.option_context("display.render_mode", "anywidget"):
-        with mock.patch.object(TableWidget, "_initial_load"):
-            widget = TableWidget(mock_df_deferred)
-            widget._initial_load_complete = True
+    with (
+        bigframes.option_context("display.render_mode", "anywidget"),
+        mock.patch.object(TableWidget, "_initial_load"),
+    ):
+        widget = TableWidget(mock_df_deferred)
+        widget._initial_load_complete = True
 
-            mock_df_deferred.to_pandas_batches.reset_mock()
+        mock_df_deferred.to_pandas_batches.reset_mock()
 
-            widget.sort_context = [{"column": "B", "ascending": False}]
+        widget.sort_context = [{"column": "B", "ascending": False}]
 
     assert mock_df_deferred.to_pandas_batches.call_count >= 1
 
@@ -510,22 +520,24 @@ def test_deferred_mode_execution_in_colab(mock_deferred_df, mock_df_deferred):
     batches.total_rows = 1
     mock_df_deferred.to_pandas_batches.return_value = batches
 
-    with mock.patch.dict(sys.modules, {"google.colab": mock.MagicMock()}):
-        with bigframes.option_context("display.render_mode", "anywidget"):
-            widget = TableWidget(mock_deferred_df)
-            widget.is_deferred_mode = True
+    with (
+        mock.patch.dict(sys.modules, {"google.colab": mock.MagicMock()}),
+        bigframes.option_context("display.render_mode", "anywidget"),
+    ):
+        widget = TableWidget(mock_deferred_df)
+        widget.is_deferred_mode = True
 
-            widget.start_execution = True
+        widget.start_execution = True
 
-            thread = getattr(widget, "_execution_thread", None)
-            if thread is not None:
-                thread.join(timeout=5)
+        thread = getattr(widget, "_execution_thread", None)
+        if thread is not None:
+            thread.join(timeout=5)
 
-            assert widget.is_deferred_mode is True
-            assert widget.table_html == ""
+        assert widget.is_deferred_mode is True
+        assert widget.table_html == ""
 
-            # Simulate frontend ping callback
-            widget.ping = 1
+        # Simulate frontend ping callback
+        widget.ping = 1
 
-            assert widget.is_deferred_mode is False
-            assert widget.table_html != ""
+        assert widget.is_deferred_mode is False
+        assert widget.table_html != ""

packages/bigframes/tests/unit/display/test_html.py

@@ -203,19 +203,20 @@ def test_repr_mimebundle_head():
     )
 
     # Mock _get_obj_metadata
-    with patch("bigframes.display.html._get_obj_metadata", return_value=(False, False)):
-        # Mock create_html_representation and create_text_representation
-        with patch(
+    with (
+        patch("bigframes.display.html._get_obj_metadata", return_value=(False, False)),
+        patch(
             "bigframes.display.html.create_html_representation", return_value="<html>"
-        ) as mock_create_html:
-            with patch(
-                "bigframes.display.plaintext.create_text_representation",
-                return_value="text",
-            ) as mock_create_text:
-                bundle = bf_html.repr_mimebundle_head(mock_df)
-
-                assert bundle == {"text/html": "<html>", "text/plain": "text"}
-                mock_df._prepare_display_df.assert_called_once()
-                mock_df._block.retrieve_repr_request_results.assert_called_once()
-                mock_create_html.assert_called_once()
-                mock_create_text.assert_called_once()
+        ) as mock_create_html,
+        patch(
+            "bigframes.display.plaintext.create_text_representation",
+            return_value="text",
+        ) as mock_create_text,
+    ):
+        bundle = bf_html.repr_mimebundle_head(mock_df)
+
+        assert bundle == {"text/html": "<html>", "text/plain": "text"}
+        mock_df._prepare_display_df.assert_called_once()
+        mock_df._block.retrieve_repr_request_results.assert_called_once()
+        mock_create_html.assert_called_once()
+        mock_create_text.assert_called_once()

packages/bigframes/tests/unit/display/test_render_mode.py

@@ -14,6 +14,8 @@
 
 import unittest.mock as mock
 
+import pytest
+
 import bigframes.display.html as bf_html
 import bigframes.pandas as bpd
 
@@ -62,7 +64,10 @@ def test_repr_mimebundle_selection_logic():
 
         # Test fallback to static deferred repr when anywidget fails
         mock_anywidget.side_effect = Exception("Anywidget failed")
-        with bpd.option_context("display.repr_mode", "deferred"):
+        with (
+            bpd.option_context("display.repr_mode", "deferred"),
+            pytest.warns(UserWarning, match="Anywidget mode is not available"),
+        ):
             bundle = bf_html.repr_mimebundle(mock_obj)
             assert bundle == {"text/plain": "deferred"}
             mock_deferred.assert_called_once()