diff --git a/pyupgrade/_plugins/unittest_aliases.py b/pyupgrade/_plugins/unittest_aliases.py index deaa4eb1..1aefc8bf 100644 --- a/pyupgrade/_plugins/unittest_aliases.py +++ b/pyupgrade/_plugins/unittest_aliases.py @@ -35,6 +35,12 @@ 'assertRaisesRegexp': 'assertRaisesRegex', } +FUNCTION_MAPPING = { + 'findTestCases': 'defaultTestLoader.loadTestsFromModule', + 'makeSuite': 'defaultTestLoader.loadTestsFromTestCase', + 'getTestCaseNames': 'defaultTestLoader.getTestCaseNames', +} + @register(ast.Call) def visit_Call( @@ -59,3 +65,15 @@ def visit_Call( new=f'self.{method_mapping[node.func.attr]}', ) yield ast_to_offset(node.func), func + elif ( + isinstance(node.func, ast.Attribute) and + isinstance(node.func.value, ast.Name) and + node.func.value.id == 'unittest' and + node.func.attr in FUNCTION_MAPPING + ): + func = functools.partial( + replace_name, + name=node.func.attr, + new=f'unittest.{FUNCTION_MAPPING[node.func.attr]}', + ) + yield ast_to_offset(node.func), func diff --git a/tests/features/unittest_aliases_test.py b/tests/features/unittest_aliases_test.py index a13c8a42..0ae33a45 100644 --- a/tests/features/unittest_aliases_test.py +++ b/tests/features/unittest_aliases_test.py @@ -63,3 +63,25 @@ def test_fix_unittest_aliases_py27(s, expected): def test_fix_unittest_aliases_py3(s, expected): ret = _fix_plugins(s, settings=Settings(min_version=(3,))) assert ret == expected + + +@pytest.mark.parametrize( + ('s', 'expected'), + ( + ( + 'unittest.findTestCases(MyTests)', + 'unittest.defaultTestLoader.loadTestsFromModule(MyTests)', + ), + ( + 'unittest.makeSuite(MyTests)', + 'unittest.defaultTestLoader.loadTestsFromTestCase(MyTests)', + ), + ( + 'unittest.getTestCaseNames(MyTests)', + 'unittest.defaultTestLoader.getTestCaseNames(MyTests)', + ), + ), +) +def test_fix_unittest_aliases_py311(s, expected): + ret = _fix_plugins(s, settings=Settings(min_version=(3, 11))) + assert ret == expected