forked from pandas-dev/pandas
/
test_rewrite_warning.py
39 lines (34 loc) · 1.12 KB
/
test_rewrite_warning.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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)