use test-specific small data instead

Author: shuoweil

tests/system/small/test_anywidget.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import traceback
-
 import pandas as pd
 import pytest
 
@@ -22,64 +20,25 @@
 pytest.importorskip("anywidget")
 
 
-@pytest.fixture(scope="module", autouse=True)
-def cleanup_session(session):
-    """Ensure comprehensive cleanup happens after all tests in this module."""
-    yield
-    try:
-        # Force cleanup of anonymous dataset and all temporary tables
-        if hasattr(session, "_anon_dataset_manager") and session._anon_dataset_manager:
-            session._anon_dataset_manager.close()
-
-        # Also call the main session cleanup
-        session.close()
-    except Exception as e:
-        traceback.print_exception(type(e), e, None)
-
-
 @pytest.fixture(scope="module")
 def paginated_pandas_df() -> pd.DataFrame:
     """Create a test DataFrame with exactly 3 pages of manually defined data."""
+    """Create a minimal test DataFrame with exactly 3 pages of 2 rows each."""
     test_data = pd.DataFrame(
         {
-            "id": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+            "id": [0, 1, 2, 3, 4, 5],
             "page_indicator": [
-                # Page 1 (rows 1-5)
+                # Page 1 (rows 1-2)
                 "page_1_row_1",
                 "page_1_row_2",
-                "page_1_row_3",
-                "page_1_row_4",
-                "page_1_row_5",
-                # Page 2 (rows 6-10)
+                # Page 2 (rows 3-4)
                 "page_2_row_1",
                 "page_2_row_2",
-                "page_2_row_3",
-                "page_2_row_4",
-                "page_2_row_5",
-                # Page 3 (rows 11-15)
+                # Page 3 (rows 5-6)
                 "page_3_row_1",
                 "page_3_row_2",
-                "page_3_row_3",
-                "page_3_row_4",
-                "page_3_row_5",
-            ],
-            "value": [
-                "data_001",
-                "data_002",
-                "data_003",
-                "data_004",
-                "data_005",
-                "data_006",
-                "data_007",
-                "data_008",
-                "data_009",
-                "data_010",
-                "data_011",
-                "data_012",
-                "data_013",
-                "data_014",
-                "data_015",
             ],
+            "value": [0, 1, 2, 3, 4, 5],
         }
     )
     return test_data
@@ -100,7 +59,7 @@ def table_widget(paginated_bf_df: bf.dataframe.DataFrame):
     """
     from bigframes.display import TableWidget
 
-    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 5):
+    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
         widget = TableWidget(paginated_bf_df)
     return widget
 
@@ -117,18 +76,14 @@ def _assert_html_matches_pandas_slice(
     """
     # Check that the unique indicator from each expected row is present.
     for _, row in expected_pd_slice.iterrows():
-        assert (
-            row["page_indicator"] in table_html
-        ), f"Expected row '{row['page_indicator']}' to be in the table HTML."
+        assert row["page_indicator"] in table_html
 
     # Create a DataFrame of all rows that should NOT be present.
     unexpected_pd_df = full_pd_df.drop(expected_pd_slice.index)
 
     # Check that no unique indicators from unexpected rows are present.
     for _, row in unexpected_pd_df.iterrows():
-        assert (
-            row["page_indicator"] not in table_html
-        ), f"Expected row '{row['page_indicator']}' NOT to be in the table HTML."
+        assert row["page_indicator"] not in table_html
 
 
 def test_repr_anywidget_initialization_set_correct_defaults(
@@ -143,21 +98,17 @@ def test_repr_anywidget_initialization_set_correct_defaults(
 
         widget = TableWidget(paginated_bf_df)
 
-        assert widget.page == 0, "Initial page should be 0."
-        assert (
-            widget.page_size == bf.options.display.max_rows
-        ), "Page size should default to max_rows option."
-        assert widget.row_count == len(
-            paginated_pandas_df
-        ), "Row count should match the source DataFrame."
+        assert widget.page == 0
+        assert widget.page_size == bf.options.display.max_rows
+        assert widget.row_count == len(paginated_pandas_df)
 
 
 def test_repr_anywidget_display_first_page_on_load(table_widget, paginated_pandas_df):
     """
     Given a widget, when it is first loaded, then it should display
     the first page of data.
     """
-    expected_slice = paginated_pandas_df.iloc[0:5]
+    expected_slice = paginated_pandas_df.iloc[0:2]
 
     html = table_widget.table_html
 
@@ -169,7 +120,7 @@ def test_repr_anywidget_navigate_to_second_page(table_widget, paginated_pandas_d
     Given a widget, when the page is set to 1, then it should display
     the second page of data.
     """
-    expected_slice = paginated_pandas_df.iloc[5:10]
+    expected_slice = paginated_pandas_df.iloc[2:4]
 
     table_widget.page = 1
     html = table_widget.table_html
@@ -183,7 +134,7 @@ def test_repr_anywidget_navigate_to_last_page(table_widget, paginated_pandas_df)
     Given a widget, when the page is set to the last page (2),
     then it should display the final page of data.
     """
-    expected_slice = paginated_pandas_df.iloc[10:15]
+    expected_slice = paginated_pandas_df.iloc[4:6]
 
     table_widget.page = 2
     html = table_widget.table_html
@@ -199,12 +150,12 @@ def test_repr_anywidget_page_clamp_to_zero_for_negative_input(
     Given a widget, when a negative page number is set,
     then the page number should be clamped to 0 and display the first page.
     """
-    expected_slice = paginated_pandas_df.iloc[0:5]
+    expected_slice = paginated_pandas_df.iloc[0:2]
 
     table_widget.page = -1
     html = table_widget.table_html
 
-    assert table_widget.page == 0, "Page should be clamped to 0."
+    assert table_widget.page == 0
     _assert_html_matches_pandas_slice(html, expected_slice, paginated_pandas_df)
 
 
@@ -215,38 +166,31 @@ def test_repr_anywidget_page_clamp_to_last_page_for_out_of_bounds_input(
     Given a widget, when a page number greater than the max is set,
     then the page number should be clamped to the last valid page.
     """
-    expected_slice = paginated_pandas_df.iloc[10:15]
+    expected_slice = paginated_pandas_df.iloc[4:6]
 
     table_widget.page = 100
     html = table_widget.table_html
 
-    assert table_widget.page == 2, "Page should be clamped to the last valid page."
+    assert table_widget.page == 2
     _assert_html_matches_pandas_slice(html, expected_slice, paginated_pandas_df)
 
 
 @pytest.mark.parametrize(
     "page, start_row, end_row",
     [
-        (0, 0, 3),
-        (1, 3, 6),
-        (2, 6, 9),
-        (3, 9, 12),
-        (4, 12, 15),
+        (0, 0, 3),  # Page 0: rows 0-2
+        (1, 3, 6),  # Page 1: rows 3-5
     ],
     ids=[
-        "Page 0 (Rows 1-3)",
-        "Page 1 (Rows 4-6)",
-        "Page 2 (Rows 7-9)",
-        "Page 3 (Rows 10-12)",
-        "Page 4 (Rows 13-15)",
+        "Page 0 (Rows 0-2)",
+        "Page 1 (Rows 3-5)",
     ],
 )
 def test_repr_anywidget_paginate_correctly_with_custom_page_size(
     paginated_bf_df, paginated_pandas_df, page, start_row, end_row
 ):
     """
     A widget should paginate correctly with a custom page size of 3.
-    This uses pytest parameterization, a strong pattern from the examples.
     """
     with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 3):
         from bigframes.display import TableWidget