TST: Mock pivot and pivot_table test for PerformanceWarning (#45112)

Author: mroeschke

pandas/tests/frame/test_stack_unstack.py

@@ -18,6 +18,7 @@
     date_range,
 )
 import pandas._testing as tm
+from pandas.core.reshape import reshape as reshape_lib
 
 
 class TestDataFrameReshape:
@@ -1826,19 +1827,26 @@ def test_unstack_unobserved_keys(self):
         tm.assert_frame_equal(recons, df)
 
     @pytest.mark.slow
-    def test_unstack_number_of_levels_larger_than_int32(self):
+    def test_unstack_number_of_levels_larger_than_int32(self, monkeypatch):
         # GH#20601
         # GH 26314: Change ValueError to PerformanceWarning
-        df = DataFrame(
-            np.random.randn(2 ** 16, 2), index=[np.arange(2 ** 16), np.arange(2 ** 16)]
-        )
-        msg = "The following operation may generate"
-        with tm.assert_produces_warning(PerformanceWarning, match=msg):
-            try:
-                df.unstack()
-            except MemoryError:
-                # Just checking the warning
-                return
+
+        class MockUnstacker(reshape_lib._Unstacker):
+            def __init__(self, *args, **kwargs):
+                # __init__ will raise the warning
+                super().__init__(*args, **kwargs)
+                raise Exception("Don't compute final result.")
+
+        with monkeypatch.context() as m:
+            m.setattr(reshape_lib, "_Unstacker", MockUnstacker)
+            df = DataFrame(
+                np.random.randn(2 ** 16, 2),
+                index=[np.arange(2 ** 16), np.arange(2 ** 16)],
+            )
+            msg = "The following operation may generate"
+            with tm.assert_produces_warning(PerformanceWarning, match=msg):
+                with pytest.raises(Exception, match="Don't compute final result."):
+                    df.unstack()
 
     def test_stack_order_with_unsorted_levels(self):
         # GH#16323

pandas/tests/reshape/test_pivot.py

@@ -23,6 +23,7 @@
 )
 import pandas._testing as tm
 from pandas.api.types import CategoricalDtype as CDT
+from pandas.core.reshape import reshape as reshape_lib
 from pandas.core.reshape.pivot import pivot_table
 
 
@@ -1991,22 +1992,27 @@ def test_pivot_string_func_vs_func(self, f, f_numpy):
         tm.assert_frame_equal(result, expected)
 
     @pytest.mark.slow
-    def test_pivot_number_of_levels_larger_than_int32(self):
+    def test_pivot_number_of_levels_larger_than_int32(self, monkeypatch):
         # GH 20601
         # GH 26314: Change ValueError to PerformanceWarning
-        df = DataFrame(
-            {"ind1": np.arange(2 ** 16), "ind2": np.arange(2 ** 16), "count": 0}
-        )
+        class MockUnstacker(reshape_lib._Unstacker):
+            def __init__(self, *args, **kwargs):
+                # __init__ will raise the warning
+                super().__init__(*args, **kwargs)
+                raise Exception("Don't compute final result.")
+
+        with monkeypatch.context() as m:
+            m.setattr(reshape_lib, "_Unstacker", MockUnstacker)
+            df = DataFrame(
+                {"ind1": np.arange(2 ** 16), "ind2": np.arange(2 ** 16), "count": 0}
+            )
 
-        msg = "The following operation may generate"
-        with tm.assert_produces_warning(PerformanceWarning, match=msg):
-            try:
-                df.pivot_table(
-                    index="ind1", columns="ind2", values="count", aggfunc="count"
-                )
-            except MemoryError:
-                # Just checking the warning
-                return
+            msg = "The following operation may generate"
+            with tm.assert_produces_warning(PerformanceWarning, match=msg):
+                with pytest.raises(Exception, match="Don't compute final result."):
+                    df.pivot_table(
+                        index="ind1", columns="ind2", values="count", aggfunc="count"
+                    )
 
     def test_pivot_table_aggfunc_dropna(self, dropna):
         # GH 22159