gh-155907: Complete PyMarshal C API tests

Add tests on PyMarshal_ReadObjectFromString() and
PyMarshal_WriteObjectToString().

Add test on PyMarshal_WriteObjectToFile(NULL).

Author: vstinner

Lib/test/test_capi/test_marshal.py

@@ -1,19 +1,54 @@
+# Test PyMarshal C API
+
 import marshal
 import os.path
+import struct
 import unittest
 
 from test import support
 from test.support import import_helper
 from test.support import os_helper
-from test.test_marshal import HelperMixin, omit_last_byte
 
 
 # Skip this test if _testcapi is are not available.
 _testcapi = import_helper.import_module('_testcapi')
 
 
+def noop_func():
+    pass
+
+NULL = None
+SIMPLE_OBJECT = 123
+# Only test a few objects: see test_marshal for more exhaustive tests
+TEST_OBJECTS = (
+    '\u20ac',
+    b'abc',
+    True,
+    123,
+    45.6,
+    7+8j,
+    'long line '*1000,
+    # Check that serializing code object is allowed (allow_code = 1)
+    noop_func.__code__,
+)
+UNMARSHALLABLE = object()
+
+# Invalid marshal data
+JUNK_BYTES = b'\xff' * 32
+
+
+def read_file(filename):
+    with open(filename, 'rb') as fp:
+        return fp.read()
+
+
+def write_file(filename, data):
+    with open(filename, 'wb') as fp:
+        fp.write(data)
+
+
 @support.cpython_only
-class CAPI_TestCase(unittest.TestCase, HelperMixin):
+class CAPI_TestCase(unittest.TestCase):
 
     def test_read_from_file_error(self):
         # A read error is reported as OSError, not EOFError.
@@ -38,89 +73,159 @@ def test_write_to_file_error(self):
             _testcapi.pymarshal_write_object_to_file(obj, '/dev/full',
                                                      marshal.version)
 
-    def test_write_unmarshallable_to_file(self):
-        self.addCleanup(os_helper.unlink, os_helper.TESTFN)
-        with self.assertRaisesRegex(ValueError, 'unmarshallable object'):
-            _testcapi.pymarshal_write_object_to_file(object(), os_helper.TESTFN,
-                                                     marshal.version)
+    def check_object(self, obj2, obj):
+        self.assertEqual(obj2, obj)
+        self.assertEqual(type(obj2), type(obj))
 
     def test_write_long_to_file(self):
-        for v in range(marshal.version + 1):
-            _testcapi.pymarshal_write_long_to_file(0x12345678, os_helper.TESTFN, v)
-            with open(os_helper.TESTFN, 'rb') as f:
-                data = f.read()
-            os_helper.unlink(os_helper.TESTFN)
-            self.assertEqual(data, b'\x78\x56\x34\x12')
+        # Test PyMarshal_WriteLongToFile()
+        write_long_to_file = _testcapi.pymarshal_write_long_to_file
+        filename = os_helper.TESTFN
+        self.addCleanup(os_helper.unlink, filename)
+
+        def mask32(value):
+            res = value & (2 ** 32  - 1)
+            if res >= 2147483648:
+                return res - 4294967296
+            else:
+                return res
+
+        limit = 2 ** 31
+        for version in range(marshal.version + 1):
+            for value in (
+                _testcapi.LONG_MIN, _testcapi.LONG_MAX,
+                -limit - 2, -limit, -limit + 2,
+                limit - 2, limit, limit + 2,
+                0, 123, -123,
+            ):
+                with self.subTest(value=value, version=version):
+                    write_long_to_file(value, filename, version)
+                    data = read_file(filename)
+                    self.assertEqual(len(data), 4)
+                    value2 = struct.unpack('<i', data)[0]
+                    self.assertEqual(value2, mask32(value))
 
     def test_write_object_to_file(self):
-        obj = ('\u20ac', b'abc', 123, 45.6, 7+8j, 'long line '*1000)
-        for v in range(marshal.version + 1):
-            _testcapi.pymarshal_write_object_to_file(obj, os_helper.TESTFN, v)
-            with open(os_helper.TESTFN, 'rb') as f:
-                data = f.read()
-            os_helper.unlink(os_helper.TESTFN)
-            self.assertEqual(marshal.loads(data), obj)
+        # Test PyMarshal_WriteObjectToFile()
+        write_object_to_file = _testcapi.pymarshal_write_object_to_file
+        filename = os_helper.TESTFN
+        self.addCleanup(os_helper.unlink, filename)
+
+        for version in range(marshal.version + 1):
+            for obj in TEST_OBJECTS:
+                with self.subTest(obj=obj, version=version):
+                    write_object_to_file(obj, filename, version)
+                    data = read_file(filename)
+                    self.assertEqual(marshal.loads(data), obj)
+
+            with self.assertRaises(SystemError):
+                write_object_to_file(NULL, filename, version)
+
+            with self.assertRaisesRegex(ValueError, 'unmarshallable object'):
+                write_object_to_file(UNMARSHALLABLE, filename, version)
 
     def test_read_short_from_file(self):
-        with open(os_helper.TESTFN, 'wb') as f:
-            f.write(b'\x34\x12xxxx')
-        r, p = _testcapi.pymarshal_read_short_from_file(os_helper.TESTFN)
-        os_helper.unlink(os_helper.TESTFN)
-        self.assertEqual(r, 0x1234)
-        self.assertEqual(p, 2)
-
-        with open(os_helper.TESTFN, 'wb') as f:
-            f.write(b'\x12')
+        # Test PyMarshal_ReadShortFromFile()
+        read_short_from_file = _testcapi.pymarshal_read_short_from_file
+        filename = os_helper.TESTFN
+        self.addCleanup(os_helper.unlink, filename)
+
+        for value in (-2**15, 2**15-1, 0, 123, -123):
+            with self.subTest(value=value):
+                data = struct.pack('<h', value) + b'xxxx'
+                write_file(filename, data)
+                value2 = read_short_from_file(filename)
+                self.assertEqual(value2, value)
+
+        write_file(filename, b'\x12')  # less than 2 bytes
         with self.assertRaises(EOFError):
-            _testcapi.pymarshal_read_short_from_file(os_helper.TESTFN)
-        os_helper.unlink(os_helper.TESTFN)
+            read_short_from_file(filename)
 
     def test_read_long_from_file(self):
-        with open(os_helper.TESTFN, 'wb') as f:
-            f.write(b'\x78\x56\x34\x12xxxx')
-        r, p = _testcapi.pymarshal_read_long_from_file(os_helper.TESTFN)
-        os_helper.unlink(os_helper.TESTFN)
-        self.assertEqual(r, 0x12345678)
-        self.assertEqual(p, 4)
-
-        with open(os_helper.TESTFN, 'wb') as f:
-            f.write(b'\x56\x34\x12')
+        # Test PyMarshal_ReadLongFromFile()
+        read_long_from_file = _testcapi.pymarshal_read_long_from_file
+        filename = os_helper.TESTFN
+        self.addCleanup(os_helper.unlink, filename)
+
+        for value in (_testcapi.INT_MIN, _testcapi.INT_MAX, 0, 123, -123):
+            with self.subTest(value=value):
+                data = struct.pack('<i', value)
+                write_file(filename, data)
+                value2 = read_long_from_file(filename)
+                self.assertEqual(value2, value)
+
+        write_file(filename, b'\x56\x34\x12')  # less than 4 bytes
+        with self.assertRaises(EOFError):
+            read_long_from_file(filename)
+
+    def check_read_object(self, read_object_func, check_pos=True):
+        filename = os_helper.TESTFN
+        self.addCleanup(os_helper.unlink, filename)
+
+        version = marshal.version
+        for obj in TEST_OBJECTS:
+            with self.subTest(obj=obj):
+                data = marshal.dumps(obj, version)
+                data += b'abc'  # following data is ignored
+                write_file(filename, data)
+                obj2, pos = read_object_func(filename)
+                self.check_object(obj2, obj)
+                if check_pos:
+                    self.assertEqual(pos, len(data))
+
+        data = marshal.dumps(SIMPLE_OBJECT, version)
+        data = data[:-1]  # truncate last byte
+        write_file(filename, data)
         with self.assertRaises(EOFError):
-            _testcapi.pymarshal_read_long_from_file(os_helper.TESTFN)
-        os_helper.unlink(os_helper.TESTFN)
+            read_object_func(filename)
+
+        write_file(filename, JUNK_BYTES)
+        with self.assertRaisesRegex(ValueError, 'bad marshal data'):
+            read_object_func(filename)
 
     def test_read_last_object_from_file(self):
-        obj = ('\u20ac', b'abc', 123, 45.6, 7+8j)
-        for v in range(marshal.version + 1):
-            data = marshal.dumps(obj, v)
-            with open(os_helper.TESTFN, 'wb') as f:
-                f.write(data + b'xxxx')
-            r, p = _testcapi.pymarshal_read_last_object_from_file(os_helper.TESTFN)
-            os_helper.unlink(os_helper.TESTFN)
-            self.assertEqual(r, obj)
-
-            with open(os_helper.TESTFN, 'wb') as f:
-                f.write(omit_last_byte(data))
-            with self.assertRaises(EOFError):
-                _testcapi.pymarshal_read_last_object_from_file(os_helper.TESTFN)
-            os_helper.unlink(os_helper.TESTFN)
+        # Test PyMarshal_ReadLastObjectFromFile()
+        read_last_object_from_file = _testcapi.pymarshal_read_last_object_from_file
+        self.check_read_object(read_last_object_from_file)
 
     def test_read_object_from_file(self):
-        obj = ('\u20ac', b'abc', 123, 45.6, 7+8j)
-        for v in range(marshal.version + 1):
-            data = marshal.dumps(obj, v)
-            with open(os_helper.TESTFN, 'wb') as f:
-                f.write(data + b'xxxx')
-            r, p = _testcapi.pymarshal_read_object_from_file(os_helper.TESTFN)
-            os_helper.unlink(os_helper.TESTFN)
-            self.assertEqual(r, obj)
-            self.assertEqual(p, len(data))
-
-            with open(os_helper.TESTFN, 'wb') as f:
-                f.write(omit_last_byte(data))
-            with self.assertRaises(EOFError):
-                _testcapi.pymarshal_read_object_from_file(os_helper.TESTFN)
-            os_helper.unlink(os_helper.TESTFN)
+        # Test PyMarshal_ReadObjectFromFile()
+        read_object_from_file = _testcapi.pymarshal_read_object_from_file
+        self.check_read_object(read_object_from_file, check_pos=False)
+
+    def test_pymarshal_readobjectfromstring(self):
+        # Test PyMarshal_ReadObjectFromString()
+        readobjectfromstring = _testcapi.pymarshal_readobjectfromstring
+        for obj in TEST_OBJECTS:
+            for version in range(marshal.version + 1):
+                with self.subTest(obj=obj, version=version):
+                    data = marshal.dumps(obj, version)
+                    obj2 = readobjectfromstring(data)
+                    self.check_object(obj2, obj)
+
+        data = marshal.dumps(SIMPLE_OBJECT, marshal.version)
+        data = data[:-1]  # truncate last byte
+        with self.assertRaises(EOFError):
+            readobjectfromstring(data)
+
+        with self.assertRaisesRegex(ValueError, 'bad marshal data'):
+            readobjectfromstring(JUNK_BYTES)
+
+    def test_pymarshal_writeobjecttostring(self):
+        # Test PyMarshal_WriteObjectToString()
+        writeobjecttostring = _testcapi.pymarshal_writeobjecttostring
+        for version in range(marshal.version + 1):
+            for obj in TEST_OBJECTS:
+                with self.subTest(obj=obj, version=version):
+                    data = writeobjecttostring(obj, version)
+                    obj2 = marshal.loads(data)
+                    self.check_object(obj2, obj)
+
+            with self.assertRaisesRegex(ValueError, 'unmarshallable object'):
+                writeobjecttostring(UNMARSHALLABLE, version)
+
+            with self.assertRaises(SystemError):
+                writeobjecttostring(NULL, version)
 
 
 if __name__ == "__main__":