3RNN/Lib/site-packages/pandas/tests/util/test_rewrite_warning.py

40 lines
1.1 KiB
Python
Raw Normal View History

2024-05-26 19:49:15 +02:00
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)