done

lib/python3.11/site-packages/pandas/tests/util/test_rewrite_warning.py

@@ -0,0 +1,39 @@
+import warnings
+
+import pytest
+
+from pandas.util._exceptions import rewrite_warning
+
+import pandas._testing as tm
+
+
+@pytest.mark.parametrize(
+    "target_category, target_message, hit",
+    [
+        (FutureWarning, "Target message", True),
+        (FutureWarning, "Target", True),
+        (FutureWarning, "get mess", True),
+        (FutureWarning, "Missed message", False),
+        (DeprecationWarning, "Target message", False),
+    ],
+)
+@pytest.mark.parametrize(
+    "new_category",
+    [
+        None,
+        DeprecationWarning,
+    ],
+)
+def test_rewrite_warning(target_category, target_message, hit, new_category):
+    new_message = "Rewritten message"
+    if hit:
+        expected_category = new_category if new_category else target_category
+        expected_message = new_message
+    else:
+        expected_category = FutureWarning
+        expected_message = "Target message"
+    with tm.assert_produces_warning(expected_category, match=expected_message):
+        with rewrite_warning(
+            target_message, target_category, new_message, new_category
+        ):
+            warnings.warn(message="Target message", category=FutureWarning)