import unittest import resource import os import subprocess import time from .resource_limits import set_resource_limits, apply_limits_to_process # Relative import for testing in same directory class ResourceLimitsTest(unittest.TestCase): def setUp(self): # Store original limits before tests self.original_cpu_limit = resource.getrlimit(resource.RLIMIT_CPU) self.original_memory_limit = resource.getrlimit(resource.RLIMIT_AS) self.original_open_files_limit = resource.getrlimit(resource.RLIMIT_NOFILE) self.original_max_processes_limit = resource.getrlimit(resource.RLIMIT_NPROC) self.original_core_dump_limit = resource.getrlimit(resource.RLIMIT_CORE) def tearDown(self): # Restore original limits after tests resource.setrlimit(resource.RLIMIT_CPU, self.original_cpu_limit) resource.setrlimit(resource.RLIMIT_AS, self.original_memory_limit) resource.setrlimit(resource.RLIMIT_NOFILE, self.original_open_files_limit) resource.setrlimit(resource.RLIMIT_NPROC, self.original_max_processes_limit) resource.setrlimit(resource.RLIMIT_CORE, self.original_core_dump_limit) def test_set_resource_limits(self): cpu_time = 10 memory_limit = 1024 * 1024 * 128 # 128MB open_files = 256 max_processes = 128 set_resource_limits(cpu_time, memory_limit, open_files, max_processes) current_cpu_limit = resource.getrlimit(resource.RLIMIT_CPU) current_memory_limit = resource.getrlimit(resource.RLIMIT_AS) current_open_files_limit = resource.getrlimit(resource.RLIMIT_NOFILE) current_max_processes_limit = resource.getrlimit(resource.RLIMIT_NPROC) self.assertEqual(current_cpu_limit[0], cpu_time) self.assertEqual(current_cpu_limit[1], cpu_time) self.assertEqual(current_memory_limit[0], memory_limit) self.assertEqual(current_memory_limit[1], memory_limit) self.assertEqual(current_open_files_limit[0], open_files) self.assertEqual(current_open_files_limit[1], open_files) self.assertEqual(current_max_processes_limit[0], max_processes) self.assertEqual(current_max_processes_limit[1], max_processes) core_dump_limit = resource.getrlimit(resource.RLIMIT_CORE) self.assertEqual(core_dump_limit, (0,0)) # Ensure core dumps are disabled. def test_cpu_limit_enforcement(self): """ Tests that the CPU limit is enforced. It spawns a subprocess that attempts to consume CPU time beyond the limit. The test checks that the subprocess terminates before its natural completion. """ cpu_time = 2 # Very short time to force termination set_resource_limits(cpu_time=cpu_time) script = """ import time start_time = time.time() while True: if time.time() - start_time > 5: # Run for 5 seconds if no limit print("Finished without CPU limit.") break pass print("Exiting...") """ try: process = subprocess.Popen(['python3', '-c', script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, preexec_fn=apply_limits_to_process) process.wait(timeout=cpu_time + 1) # Wait for the process with a small grace period. return_code = process.returncode stdout, stderr = process.communicate() self.assertNotEqual(return_code, 0, "Process should have been terminated by CPU limit.") self.assertIn("Exiting...", stdout.decode(), "Process should have printed 'Exiting...'") print(f"Stdout: {stdout.decode()}") print(f"Stderr: {stderr.decode()}") except subprocess.TimeoutExpired: self.fail("Process did not terminate within the expected time.") def test_memory_limit_enforcement(self): """Tests that memory limit is enforced, causing OOM.""" memory_limit = 64 * 1024 * 1024 # 64 MB set_resource_limits(memory_limit=memory_limit) script = f""" import time try: large_list = bytearray({memory_limit * 2}) # Try to allocate twice the memory print("Large list allocated successfully (which is unexpected).") except MemoryError as e: print(f"Memory allocation failed as expected: {{e}}") except Exception as e: print(f"Unexpected exception: {{e}}") """ process = subprocess.Popen(['python3', '-c', script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, preexec_fn=apply_limits_to_process) stdout, stderr = process.communicate(timeout=10) return_code = process.returncode stdout_str = stdout.decode() stderr_str = stderr.decode() print(f"Stdout: {stdout_str}") print(f"Stderr: {stderr_str}") self.assertTrue("Memory allocation failed as expected" in stdout_str or "OSError: [Errno 12] Cannot allocate memory" in stderr_str or "Killed" in stderr_str, "MemoryError should have been raised or process should have been killed") def test_open_files_limit_enforcement(self): open_files_limit = 10 set_resource_limits(open_files=open_files_limit) script = f""" import os try: files = [] for i in range({open_files_limit + 5}): # Exceed the limit files.append(open(f"test_file_{{i}}.txt", "w")) files[-1].write("Test data") # Write something to make sure file is really open print("Opened all files successfully (which is unexpected).") except OSError as e: print(f"Failed to open files due to limit: {{e}}") finally: for f in files: try: f.close() except: pass # Ignore errors on close for i in range({open_files_limit + 5}): try: os.remove(f"test_file_{{i}}.txt") except: pass # Ignore errors on deleting, file might have been never created """ process = subprocess.Popen(['python3', '-c', script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, preexec_fn=apply_limits_to_process) stdout, stderr = process.communicate(timeout=10) return_code = process.returncode stdout_str = stdout.decode() stderr_str = stderr.decode() print(f"Stdout: {stdout_str}") print(f"Stderr: {stderr_str}") self.assertTrue("Failed to open files due to limit" in stdout_str, "OSError due to open file limit should have been raised.") if __name__ == '__main__': unittest.main()