Add test on --no-compat option

vstinner · vstinner · commit 1ba32a448f06 · 2021-04-14T14:41:25.000+02:00
diff --git a/tests/test_upgrade_pythoncapi.py b/tests/test_upgrade_pythoncapi.py
@@ -11,8 +11,10 @@
 import upgrade_pythoncapi   # noqa
 
 
-def patch(source):
+def patch(source, no_compat=False):
     args = ['script', 'mod.c']
+    if no_compat:
+        args.append('--no-compat')
     patcher = upgrade_pythoncapi.Patcher(args)
     return patcher.patch(source)
 
@@ -99,10 +101,10 @@ def test_patch_directory(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             self._test_patch_file(tmp_dir)
 
-    def check_replace(self, source, expected):
+    def check_replace(self, source, expected, **kwargs):
         source = reformat(source)
         expected = reformat(expected)
-        self.assertEqual(patch(source), expected)
+        self.assertEqual(patch(source, **kwargs), expected)
 
     def check_dont_replace(self, source):
         source = reformat(source)
@@ -537,5 +539,21 @@ def test_py_is(self):
         """)
 
 
+    def test_no_compat(self):
+        # Don't add "#include "pythoncapi_compat.h"
+        source = """
+            void test_type(PyObject *obj, PyTypeObject *type)
+            {
+                obj->ob_type = type;
+            }
+        """
+        expected = """
+            void test_type(PyObject *obj, PyTypeObject *type)
+            {
+                Py_SET_TYPE(obj, type);
+            }
+        """
+        self.check_replace(source, expected, no_compat=True)
+
 if __name__ == "__main__":
     unittest.main()
