diff --git a/tests/dangerously_run_function_in_subprocess.py b/tests/dangerously_run_function_in_subprocess.py new file mode 100644 index 0000000000..e7c068d6d9 --- /dev/null +++ b/tests/dangerously_run_function_in_subprocess.py @@ -0,0 +1,46 @@ +import inspect +import subprocess +import sys +import textwrap +from typing import Any, Callable + + +def dangerously_run_function_in_subprocess(func: Callable[[], Any]) -> tuple[str, str, int]: + """**Use with caution! This should _only_ be used with trusted code!** + + Extracts a function's source and runs it in a separate subprocess. Returns stdout, stderr, and return code + from the subprocess. + + This is useful for tests where an isolated environment is required. + + The function to be called must not have any arguments and must not have any closures over the scope in which is was + defined. + + Any modules that the function depends on must be imported inside the function. + """ + + source_code = inspect.getsource(func) + + # Must dedent the source code to avoid indentation errors + dedented_source_code = textwrap.dedent(source_code) + + # Get the function name so we can call it in the subprocess + func_name = func.__name__ + + # Create a script that calls the function + script = f""" +import sys + +{dedented_source_code} + +if __name__ == "__main__": + {func_name}() +""" + + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], # Run the script in a subprocess + capture_output=True, # Capture stdout and stderr + text=True, + ) + + return result.stdout, result.stderr, result.returncode diff --git a/tests/test_dangerously_run_function_in_subprocess.py b/tests/test_dangerously_run_function_in_subprocess.py new file mode 100644 index 0000000000..a1e21cc507 --- /dev/null +++ b/tests/test_dangerously_run_function_in_subprocess.py @@ -0,0 +1,57 @@ +from tests.dangerously_run_function_in_subprocess import dangerously_run_function_in_subprocess + + +def test_simple_function(): + def test_func(): + print("Hello, Test!") + + stdout, stderr, returncode = dangerously_run_function_in_subprocess(test_func) + + assert returncode == 0 + assert stdout.strip() == "Hello, Test!" + assert stderr == "" + + +def test_function_with_error(): + def test_func(): + raise ValueError("This is an error") + + _stdout, stderr, returncode = dangerously_run_function_in_subprocess(test_func) + + assert returncode != 0 # Should fail + assert "ValueError: This is an error" in stderr + + +def test_function_with_imports(): + def test_func(): + import math + + print(math.sqrt(4)) + + stdout, stderr, returncode = dangerously_run_function_in_subprocess(test_func) + + assert returncode == 0 + assert stdout.strip() == "2.0" + assert stderr == "" + + +def test_function_with_sys_exit(): + def test_func(): + import sys + + sys.exit(42) + + _stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func) + + assert returncode == 42 # Should return the custom exit code + + +def test_function_with_closure(): + foo = "bar" + + def test_func(): + print(foo) + + _stdout, _stderr, returncode = dangerously_run_function_in_subprocess(test_func) + + assert returncode == 1 # Should fail because of closure