diff --git a/tests/test_rencrypt.py b/tests/test_rencrypt.py index 7a80db5..fd4e0fb 100644 --- a/tests/test_rencrypt.py +++ b/tests/test_rencrypt.py @@ -9,6 +9,39 @@ import unittest import numpy as np from zeroize import zeroize1, mlock, munlock +import platform + + +def setup_memory_limit(): + if not platform.system() == "Windows": + return + + import ctypes + from ctypes import wintypes + + # Define the Windows API functions + kernel32 = ctypes.WinDLL('kernel32', use_last_error=True) + + GetCurrentProcess = kernel32.GetCurrentProcess + GetCurrentProcess.restype = wintypes.HANDLE + + SetProcessWorkingSetSize = kernel32.SetProcessWorkingSetSize + SetProcessWorkingSetSize.restype = wintypes.BOOL + SetProcessWorkingSetSize.argtypes = [wintypes.HANDLE, ctypes.c_size_t, ctypes.c_size_t] + + # Get the handle of the current process + current_process = GetCurrentProcess() + + # Set the working set size + min_size = 6 * 1024 * 1024 # Minimum working set size + max_size = 10 * 1024 * 1024 # Maximum working set size + + result = SetProcessWorkingSetSize(current_process, min_size, max_size) + + if not result: + error_code = ctypes.get_last_error() + error_message = ctypes.FormatError(error_code) + raise RuntimeError(f"SetProcessWorkingSetSize failed with error code {error_code}: {error_message}") def read_file_in_chunks(file_path, buf): @@ -437,6 +470,10 @@ def seal_and_open_in_place_file(this, cipher_meta): class TestStringMethods(unittest.TestCase): + + @classmethod + def setUpClass(cls): + setup_memory_limit() def test_encrypt(self): seal_and_open_in_place(self, CipherMeta.Ring(RingAlgorithm.ChaCha20Poly1305))