tests

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))