feat: rewrite deprecated unittest methods

asottile · Jun 6, 2022 · 448f49c · 448f49c
1 parent 09ad799
commit 448f49c
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git 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
@@ -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