Skip to content

Commit 448f49c

Browse files
committed
feat: rewrite deprecated unittest methods
1 parent 09ad799 commit 448f49c

File tree

2 files changed

+40
-0
lines changed

2 files changed

+40
-0
lines changed

pyupgrade/_plugins/unittest_aliases.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,12 @@
3535
'assertRaisesRegexp': 'assertRaisesRegex',
3636
}
3737

38+
FUNCTION_MAPPING = {
39+
'findTestCases': 'defaultTestLoader.loadTestsFromModule',
40+
'makeSuite': 'defaultTestLoader.loadTestsFromTestCase',
41+
'getTestCaseNames': 'defaultTestLoader.getTestCaseNames',
42+
}
43+
3844

3945
@register(ast.Call)
4046
def visit_Call(
@@ -59,3 +65,15 @@ def visit_Call(
5965
new=f'self.{method_mapping[node.func.attr]}',
6066
)
6167
yield ast_to_offset(node.func), func
68+
elif (
69+
isinstance(node.func, ast.Attribute) and
70+
isinstance(node.func.value, ast.Name) and
71+
node.func.value.id == 'unittest' and
72+
node.func.attr in FUNCTION_MAPPING
73+
):
74+
func = functools.partial(
75+
replace_name,
76+
name=node.func.attr,
77+
new=f'unittest.{FUNCTION_MAPPING[node.func.attr]}',
78+
)
79+
yield ast_to_offset(node.func), func

tests/features/unittest_aliases_test.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,3 +63,25 @@ def test_fix_unittest_aliases_py27(s, expected):
6363
def test_fix_unittest_aliases_py3(s, expected):
6464
ret = _fix_plugins(s, settings=Settings(min_version=(3,)))
6565
assert ret == expected
66+
67+
68+
@pytest.mark.parametrize(
69+
('s', 'expected'),
70+
(
71+
(
72+
'unittest.findTestCases(MyTests)',
73+
'unittest.defaultTestLoader.loadTestsFromModule(MyTests)',
74+
),
75+
(
76+
'unittest.makeSuite(MyTests)',
77+
'unittest.defaultTestLoader.loadTestsFromTestCase(MyTests)',
78+
),
79+
(
80+
'unittest.getTestCaseNames(MyTests)',
81+
'unittest.defaultTestLoader.getTestCaseNames(MyTests)',
82+
),
83+
),
84+
)
85+
def test_fix_unittest_aliases_py311(s, expected):
86+
ret = _fix_plugins(s, settings=Settings(min_version=(3, 11)))
87+
assert ret == expected

0 commit comments

Comments
 (0)