Improve the assertion checker: more precise messages, and more of them

edx_lint/pylint/right_assert_check.py

@@ -28,6 +28,40 @@ class AssertChecker(BaseChecker):
 
     AFFECTED_ASSERTS = ["assertTrue", "assertFalse"]
 
+    BETTER_COMPARES = {
+        "==": "assertEqual",
+        "!=": "assertNotEqual",
+        "in": "assertIn",
+        "not in": "assertNotIn",
+        "<": "assertLess",
+        "<=": "assertLessEqual",
+        ">": "assertGreater",
+        ">=": "assertGreaterEqual",
+        "is": "assertIs",
+        "is not": "assertIsNot",
+    }
+
+    BETTER_NONE_COMPARES = {
+        "==": "assertIsNone",
+        "is": "assertIsNone",
+        "!=": "assertIsNotNone",
+        "is not": "assertIsNotNone",
+    }
+
+    INVERTED_PAIRS = [
+        ("assertEqual", "assertNotEqual"),
+        ("assertIn", "assertNotIn"),
+        ("assertLess", "assertGreaterEqual"),
+        ("assertGreater", "assertLessEqual"),
+        ("assertIs", "assertIsNot"),
+        ("assertIsNone", "assertIsNotNone"),
+    ]
+
+    INVERTED = {}
+    for yup, nope in INVERTED_PAIRS:
+        INVERTED[yup] = nope
+        INVERTED[nope] = yup
+
     MESSAGE_ID = 'wrong-assert-type'
     msgs = {
         'C%d90' % BASE_ID: (
@@ -51,39 +85,21 @@ class AssertChecker(BaseChecker):
             return
 
         first_arg = node.args[0]
-        if not isinstance(first_arg, astroid.Compare):
-            # Not a comparison, so this is probably ok:
-            return
-
-        if first_arg.ops[0][0] in ["==", "!="]:
-            # An assertTrue/False with a compare should be assertEqual:
-            self.add_message(
-                self.MESSAGE_ID,
-                args="%s(%s) should be assertEqual or assertNotEqual" % (
-                    node.func.attrname,
-                    first_arg.as_string(),
-                ),
-                node=node,
-            )
-
-        elif first_arg.ops[0][0] in ["in", "not in"]:
-            # An assertTrue/False with an in statement should be assertIn:
-            self.add_message(
-                self.MESSAGE_ID,
-                args="%s(%s) should be assertIn or assertNotIn" % (
-                    node.func.attrname,
-                    first_arg.as_string(),
-                ),
-                node=node,
-            )
-
-        elif "<" in first_arg.ops[0][0] or ">" in first_arg.ops[0][0]:
-            # An assertTrue/False with a comparison should be assertGreater or assertLess:
+        existing_code = "%s(%s)" % (node.func.attrname, first_arg.as_string())
+
+        if isinstance(first_arg, astroid.Compare):
+            compare = first_arg.ops[0][0]
+            right = first_arg.ops[0][1]
+            if isinstance(right, astroid.Const) and right.value is None:
+                # Comparing to None, handle specially.
+                better = self.BETTER_NONE_COMPARES[compare]
+            else:
+                better = self.BETTER_COMPARES[compare]
+
+            if node.func.attrname == "assertFalse":
+                better = self.INVERTED[better]
             self.add_message(
                 self.MESSAGE_ID,
-                args="%s(%s) should be assertGreater or assertLess" % (
-                    node.func.attrname,
-                    first_arg.as_string(),
-                ),
+                args="%s should be %s" % (existing_code, better),
                 node=node,
             )

test/input/func_asserts_check.py

@@ -37,3 +37,22 @@ class TestStringMethods(unittest.TestCase):
 
         self.assertTrue(1 > 0)
         self.assertFalse(1 < 2)
+
+        my_zero = 0
+
+        self.assertTrue(my_zero is 0)
+        self.assertFalse(my_zero is 1)
+        self.assertTrue(my_zero is not 1)
+        self.assertFalse(my_zero is not 0)
+
+        my_none = None
+
+        self.assertTrue(my_none is None)
+        self.assertFalse(my_zero is None)
+        self.assertTrue(my_zero != None)
+
+    def test_wrong_but_with_pragma(self):
+        """
+        This uses the wrong assert, but has a pragma to quiet the message.
+        """
+        self.assertTrue("a" in "lala")      # pylint: disable=wrong-assert-type

test/messages/func_asserts_check.txt

-C: 32:TestStringMethods.test_wrong_usage: assertTrue('foo'.upper() == 'FOO') should be assertEqual or assertNotEqual
-C: 33:TestStringMethods.test_wrong_usage: assertFalse(500 == 501) should be assertEqual or assertNotEqual
-C: 35:TestStringMethods.test_wrong_usage: assertTrue('a' in 'lala') should be assertIn or assertNotIn
-C: 36:TestStringMethods.test_wrong_usage: assertFalse('b' not in 'lala') should be assertIn or assertNotIn
-C: 38:TestStringMethods.test_wrong_usage: assertTrue(1 > 0) should be assertGreater or assertLess
-C: 39:TestStringMethods.test_wrong_usage: assertFalse(1 < 2) should be assertGreater or assertLess
\ No newline at end of file
+C: 32:TestStringMethods.test_wrong_usage: assertTrue('foo'.upper() == 'FOO') should be assertEqual
+C: 33:TestStringMethods.test_wrong_usage: assertFalse(500 == 501) should be assertNotEqual
+C: 35:TestStringMethods.test_wrong_usage: assertTrue('a' in 'lala') should be assertIn
+C: 36:TestStringMethods.test_wrong_usage: assertFalse('b' not in 'lala') should be assertIn
+C: 38:TestStringMethods.test_wrong_usage: assertTrue(1 > 0) should be assertGreater
+C: 39:TestStringMethods.test_wrong_usage: assertFalse(1 < 2) should be assertGreaterEqual
+C: 43:TestStringMethods.test_wrong_usage: assertTrue(my_zero is 0) should be assertIs
+C: 44:TestStringMethods.test_wrong_usage: assertFalse(my_zero is 1) should be assertIsNot
+C: 45:TestStringMethods.test_wrong_usage: assertTrue(my_zero is not 1) should be assertIsNot
+C: 46:TestStringMethods.test_wrong_usage: assertFalse(my_zero is not 0) should be assertIs
+C: 50:TestStringMethods.test_wrong_usage: assertTrue(my_none is None) should be assertIsNone
+C: 51:TestStringMethods.test_wrong_usage: assertFalse(my_zero is None) should be assertIsNotNone
+C: 52:TestStringMethods.test_wrong_usage: assertTrue(my_zero != None) should be assertIsNotNone