1 files changed, 49 insertions, 3 deletions
diff --git a/Lib/test/test_memoryio.py b/Lib/test/test_memoryio.py
index 9ef293e708..23be5e9bb6 100644
--- a/Lib/test/test_memoryio.py
+++ b/Lib/test/test_memoryio.py
@@ -9,6 +9,7 @@ from test import support
 import io
 import _pyio as pyio
 import pickle
+import sys
 
 class MemorySeekTestMixin:
 
@@ -711,12 +712,57 @@ class CBytesIOTest(PyBytesIOTest):
 
     @support.cpython_only
     def test_sizeof(self):
-        basesize = support.calcobjsize('P2nN2Pn')
+        basesize = support.calcobjsize('P2nN2PnP')
         check = self.check_sizeof
         self.assertEqual(object.__sizeof__(io.BytesIO()), basesize)
         check(io.BytesIO(), basesize )
-        check(io.BytesIO(b'a'), basesize + 1 + 1 )
-        check(io.BytesIO(b'a' * 1000), basesize + 1000 + 1 )
+        check(io.BytesIO(b'a'), basesize + 1 )
+        check(io.BytesIO(b'a' * 1000), basesize + 1000)
+
+    # Various tests of copy-on-write behaviour for BytesIO.
+
+    def _test_cow_mutation(self, mutation):
+        # Common code for all BytesIO copy-on-write mutation tests.
+        imm = b' ' * 1024
+        old_rc = sys.getrefcount(imm)
+        memio = self.ioclass(imm)
+        self.assertEqual(sys.getrefcount(imm), old_rc + 1)
+        mutation(memio)
+        self.assertEqual(sys.getrefcount(imm), old_rc)
+
+    @support.cpython_only
+    def test_cow_truncate(self):
+        # Ensure truncate causes a copy.
+        def mutation(memio):
+            memio.truncate(1)
+        self._test_cow_mutation(mutation)
+
+    @support.cpython_only
+    def test_cow_write(self):
+        # Ensure write that would not cause a resize still results in a copy.
+        def mutation(memio):
+            memio.seek(0)
+            memio.write(b'foo')
+        self._test_cow_mutation(mutation)
+
+    @support.cpython_only
+    def test_cow_setstate(self):
+        # __setstate__ should cause buffer to be released.
+        memio = self.ioclass(b'foooooo')
+        state = memio.__getstate__()
+        def mutation(memio):
+            memio.__setstate__(state)
+        self._test_cow_mutation(mutation)
+
+    @support.cpython_only
+    def test_cow_mutable(self):
+        # BytesIO should accept only Bytes for copy-on-write sharing, since
+        # arbitrary buffer-exporting objects like bytearray() aren't guaranteed
+        # to be immutable.
+        ba = bytearray(1024)
+        old_rc = sys.getrefcount(ba)
+        memio = self.ioclass(ba)
+        self.assertEqual(sys.getrefcount(ba), old_rc)
 
 class CStringIOTest(PyStringIOTest):
     ioclass = io.StringIO