diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py index 4d09bd7573d0e6..03f1a068b2e922 100644 --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -3455,6 +3455,22 @@ def built_with_c_assertions(): return bool(_testlimitedcapi._py_getbuiltwithassert()) +def inject_memory_error(start=0, stop=0): + """ + Memory allocation fails after 'start' allocation requests, and until 'stop' + allocation requests except when 'stop' is negative or equal to 0 (default) + in which case allocation failures never stop. + + Raise SkipTest if the _testcapi extension module is missing + """ + try: + import _testcapi + except ImportError: + raise unittest.SkipTest("_testcapi required") + + _testcapi.set_nomemory(start, stop) + + @contextlib.contextmanager def inject_memory_error_cm(start=0, stop=0): """ diff --git a/Lib/test/test_atexit.py b/Lib/test/test_atexit.py index 33c37648da31fc..33d3eb93394781 100644 --- a/Lib/test/test_atexit.py +++ b/Lib/test/test_atexit.py @@ -197,14 +197,14 @@ def test_atexit_with_low_memory(self): # callback doesn't cause an infinite loop during finalization. code = textwrap.dedent(""" import atexit - import _testcapi + from test.support import inject_memory_error def callback(): print("hello") atexit.register(callback) # Simulate low memory condition - _testcapi.set_nomemory(0) + inject_memory_error() """) with os_helper.temp_dir() as temp_dir: diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py index 2e3383d73bf2e6..4b4857295663dc 100644 --- a/Lib/test/test_bytes.py +++ b/Lib/test/test_bytes.py @@ -51,16 +51,10 @@ def __index__(self): @contextlib.contextmanager -def inject_memory_error(testcase, start): - # Raise SkipTest if _testcapi extension module is missing - _testcapi = import_helper.import_module('_testcapi') - +def inject_memory_error(testcase, start=0): with testcase.assertRaises(MemoryError): - try: - _testcapi.set_nomemory(start) + with support.inject_memory_error_cm(start): yield - finally: - _testcapi.remove_mem_hooks() class BaseBytesTest: @@ -1607,7 +1601,7 @@ def test_resize_error(self): del ba[:offset] else: expected = ba.copy() - with inject_memory_error(self, 0): + with inject_memory_error(self): ba.resize(1024) self.assertEqual(ba, expected) @@ -1618,7 +1612,7 @@ def test_resize_error(self): del ba[:offset] else: expected = ba.copy() - with inject_memory_error(self, 0): + with inject_memory_error(self): ba.resize(1) self.assertEqual(ba, expected) @@ -1691,21 +1685,20 @@ def test_take_bytes_error(self): # gh-157242: If bytearray.take_bytes() fails (MemoryError), # the bytearray must be left unchanged. - for logical_offset, to_take, mem_errors in ( + for logical_offset, to_take, start_list in ( (True, 5, (0, 1)), (False, 5, (0, 1)), (True, None, (0,)), ): - for mem_error in mem_errors: - with self.subTest(logical_offset=logical_offset, - to_take=to_take, mem_error=mem_error): + for start in start_list: + with self.subTest(logical_offset=logical_offset, start=start): ba = bytearray(b'0123456789') if logical_offset: expected = ba[3:] del ba[:3] else: expected = ba.copy() - with inject_memory_error(self, mem_error): + with inject_memory_error(self, start): ba.take_bytes(to_take) self.assertEqual(ba, expected) diff --git a/Lib/test/test_capi/test_bytes.py b/Lib/test/test_capi/test_bytes.py index 05c7f90eec7cbc..ee66d974bd3014 100644 --- a/Lib/test/test_capi/test_bytes.py +++ b/Lib/test/test_capi/test_bytes.py @@ -455,6 +455,23 @@ def test_resize(self): writer.resize(_testcapi.PY_SSIZE_T_MAX) self.assertEqual(writer.finish(), b'x' * size) + @support.nomemtest + def test_resize_error(self): + # Test PyBytesWriter_Resize() error + init = b'x' * self.LARGE_BUFFER + writer = self.create_writer(len(init)) + writer.write(0, init) + size = len(init) + 100 + with self.assertRaises(MemoryError): + with support.inject_memory_error_cm(): + writer.resize(size) + suffix = b'still working' + writer.write_bytes(suffix, -1) + self.assertEqual(writer.finish(), init + suffix) + + # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize) + # if the new size is smaller than the allocated size + def test_grow(self): # Test PyBytesWriter_Grow() writer = self.create_writer(0) @@ -513,24 +530,20 @@ def test_grow(self): self.assertEqual(writer.finish(), b'x' * size) @support.nomemtest - def test_resize_error(self): - # Test PyBytesWriter_Resize() error + def test_grow_error(self): + # Test PyBytesWriter_Grow() error init = b'x' * self.LARGE_BUFFER writer = self.create_writer(len(init)) writer.write(0, init) - size = len(init) + 100 - try: - with self.assertRaises(MemoryError): - _testcapi.set_nomemory(0) - writer.resize(size) - finally: - _testcapi.remove_mem_hooks() + with self.assertRaises(MemoryError): + with support.inject_memory_error_cm(): + writer.grow(100) suffix = b'still working' writer.write_bytes(suffix, -1) self.bytes_equal(writer.finish(), init + suffix) - # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize) - # if the new size is smaller than the allocated size + # Note: PyBytesWriter_Grow() leaves the buffer unchanged (no resize) + # if grow is negative. def test_format_i(self): # Test PyBytesWriter_Format() @@ -554,11 +567,8 @@ def test_format_s_memory_error(self): writer = self.create_writer() s = b'x' * 500 with self.assertRaises(MemoryError): - try: - _testcapi.set_nomemory(0, 0) + with support.inject_memory_error_cm(): writer.format_s(b's=%s', s) - finally: - _testcapi.remove_mem_hooks() writer.write_bytes(b'after', -1) @threading_helper.requires_working_threading() diff --git a/Lib/test/test_class.py b/Lib/test/test_class.py index 7bd6d966e8536b..3fdc87700f99ad 100644 --- a/Lib/test/test_class.py +++ b/Lib/test/test_class.py @@ -1016,8 +1016,6 @@ class C: @support.nomemtest @isolation.runInSubprocess() def test_detach_materialized_dict_no_memory(self): - import _testcapi - class A: def __init__(self): self.a = 1 @@ -1032,11 +1030,8 @@ def __init__(self): d = a.__dict__ try: with support.catch_unraisable_exception() as ex: - _testcapi.set_nomemory(n, n + 1) - try: + with support.inject_memory_error_cm(n, n + 1): del a - finally: - _testcapi.remove_mem_hooks() exc_type = ex.unraisable and ex.unraisable.exc_type except MemoryError: # The failing allocation was not in the deallocation code. diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py index 4fc851c0ce464b..adce4ee0cb60d8 100644 --- a/Lib/test/test_dict.py +++ b/Lib/test/test_dict.py @@ -461,15 +461,11 @@ def test_copy_noncompact(self): @support.nomemtest def test_copy_empty_no_memory(self): - import _testcapi d = {} d_copy = d.copy with self.assertRaises(MemoryError): - try: - _testcapi.set_nomemory(0) + with support.inject_memory_error_cm(): d_copy() - finally: - _testcapi.remove_mem_hooks() def test_get(self): d = {} diff --git a/Lib/test/test_exceptions.py b/Lib/test/test_exceptions.py index 992b9a7a8ae8bb..c6e31b65809b1d 100644 --- a/Lib/test/test_exceptions.py +++ b/Lib/test/test_exceptions.py @@ -1663,14 +1663,14 @@ def test_recursion_normalizing_with_no_memory(self): # the size of the list of preallocated MemoryError instances, the # Fatal Python error message mentions MemoryError. code = """if 1: - import _testcapi + from test import support class C(): pass def recurse(cnt): cnt -= 1 if cnt: recurse(cnt) else: - _testcapi.set_nomemory(0) + support.inject_memory_error() C() recurse(16) """ @@ -1846,9 +1846,10 @@ def test_unhandled(self): @support.nomemtest def test_memory_error_in_PyErr_PrintEx(self): code = """if 1: - import _testcapi + from test import support + stop = %d class C(): pass - _testcapi.set_nomemory(0, %d) + support.inject_memory_error(0, stop) C() """ @@ -2013,8 +2014,8 @@ def test_exec_set_nomemory_hang(self): warmup_code = "a = list(range(0, 1))\n" * 60 user_input = warmup_code + dedent(""" try: - import _testcapi - _testcapi.set_nomemory(0) + from test import support + support.inject_memory_error() b = list(range(1000, 2000)) except Exception as e: import traceback diff --git a/Lib/test/test_interpreters/test_stress.py b/Lib/test/test_interpreters/test_stress.py index 50d2444a4c72d3..56bfc12f951f79 100644 --- a/Lib/test/test_interpreters/test_stress.py +++ b/Lib/test/test_interpreters/test_stress.py @@ -78,15 +78,9 @@ def run(): @support.nomemtest def test_create_interpreter_no_memory(self): - import _testcapi - - assertion = self.assertRaises(InterpreterError) - try: - _testcapi.set_nomemory(0, 1) - with assertion: + with self.assertRaises(InterpreterError): + with support.inject_memory_error_cm(0, 1): _interpreters.create() - finally: - _testcapi.remove_mem_hooks() if __name__ == '__main__': diff --git a/Lib/test/test_io/test_memoryio.py b/Lib/test/test_io/test_memoryio.py index 6e436b5a547836..7f2eddd7ae7839 100644 --- a/Lib/test/test_io/test_memoryio.py +++ b/Lib/test/test_io/test_memoryio.py @@ -5,7 +5,6 @@ import unittest from test import support -from test.support import import_helper import gc import io @@ -654,18 +653,14 @@ def __buffer__(self, flags): def test_memory_error(self): # gh-157242: io.BytesIO() must not close the file on MemoryError - _testcapi = import_helper.import_module('_testcapi') # write() stream = self.ioclass() stream.write(self.buftype('abc')) + data = self.buftype('def') with self.assertRaises(MemoryError): - try: - data = self.buftype('def') - _testcapi.set_nomemory(0) + with support.inject_memory_error_cm(): stream.write(data) - finally: - _testcapi.remove_mem_hooks() stream.write(self.buftype('123')) self.assertEqual(stream.getvalue(), self.buftype('abc123')) @@ -674,11 +669,8 @@ def test_memory_error(self): stream = self.ioclass() stream.write(data) with self.assertRaises(MemoryError): - try: - _testcapi.set_nomemory(0) + with support.inject_memory_error_cm(): stream.truncate(5) - finally: - _testcapi.remove_mem_hooks() self.assertEqual(stream.getvalue(), data) diff --git a/Lib/test/test_list.py b/Lib/test/test_list.py index 44ecd62d6735de..a0a1259ef6810c 100644 --- a/Lib/test/test_list.py +++ b/Lib/test/test_list.py @@ -329,11 +329,12 @@ def test_tier2_invalidates_iterator(self): def test_no_memory(self): # gh-118331: Make sure we don't crash if list allocation fails code = textwrap.dedent(""" - import _testcapi, sys + from test import support + import sys # Prime the freelist l = [None] del l - _testcapi.set_nomemory(0) + support.inject_memory_error() l = [None] """) rc, _, _ = assert_python_failure("-c", code) diff --git a/Lib/test/test_pyexpat.py b/Lib/test/test_pyexpat.py index ef4d960aa4963e..d1aa0a8f868eaa 100644 --- a/Lib/test/test_pyexpat.py +++ b/Lib/test/test_pyexpat.py @@ -1028,18 +1028,9 @@ def test_error_path_no_crash(self): parser.buffer_text = True rc_before = sys.getrefcount(parser) - # We avoid self.assertRaises(MemoryError) here because the - # context manager itself needs memory allocations that fail - # while the nomemory hook is active. - self.testcapi.set_nomemory(1, 10) - raised = False - try: - parser.ExternalEntityParserCreate(None) - except MemoryError: - raised = True - finally: - self.testcapi.remove_mem_hooks() - self.assertTrue(raised, "MemoryError not raised") + with self.assertRaises(MemoryError): + with support.inject_memory_error_cm(1, 10): + parser.ExternalEntityParserCreate(None) rc_after = sys.getrefcount(parser) self.assertEqual(rc_after, rc_before) diff --git a/Lib/test/test_repl.py b/Lib/test/test_repl.py index 0f5009c3212490..96c3d4c4de692d 100644 --- a/Lib/test/test_repl.py +++ b/Lib/test/test_repl.py @@ -104,10 +104,11 @@ def test_no_memory(self): # no memory. Check also that the fix does not break the interactive # loop when an exception is raised. user_input = """ - import sys, _testcapi + import sys + from test import support 1/0 print('After the exception.') - _testcapi.set_nomemory(0) + support.inject_memory_error() sys.exit(0) """ user_input = dedent(user_input) diff --git a/Lib/test/test_str.py b/Lib/test/test_str.py index 979bfe36fff680..bef0e62fb42bcc 100644 --- a/Lib/test/test_str.py +++ b/Lib/test/test_str.py @@ -613,14 +613,9 @@ def test_replace_oom(self): s1 = "轘" * 4 s2 = "&" s3 = "&" - assertion = self.assertRaises(MemoryError) - _testcapi.set_nomemory(0, 0) - try: - # No allocations made in the test itself: - with assertion: + with self.assertRaises(MemoryError): + with support.inject_memory_error_cm(): s1.replace(s2, s3) # this line used to crash before - finally: - _testcapi.remove_mem_hooks() def test_repeat_id_preserving(self): a = '123abc1@' diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py index 66657460ec4926..ccbb0d3f179053 100644 --- a/Lib/test/test_weakref.py +++ b/Lib/test/test_weakref.py @@ -1035,8 +1035,8 @@ def test_no_memory_when_clearing(self): # gh-118331: Make sure we do not raise an exception from the destructor # when clearing weakrefs if allocating the intermediate tuple fails. code = textwrap.dedent(""" - import _testcapi import weakref + from test import support class TestObj: pass @@ -1048,7 +1048,7 @@ def callback(obj): # The choice of 50 is arbitrary, but must be large enough to ensure # the allocation won't be serviced by the free list. wrs = [weakref.ref(obj, callback) for _ in range(50)] - _testcapi.set_nomemory(0) + support.inject_memory_error() del obj """).strip() res, _ = script_helper.run_python_until_end("-c", code)