Make security tests more explicit

huggingface · Mar 4, 2025 · e6ed3f3 · e6ed3f3
1 parent a39148d
commit e6ed3f3
Showing 1 changed file with 56 additions and 28 deletions.
diff --git a/tests/test_local_python_executor.py b/tests/test_local_python_executor.py
@@ -1421,30 +1421,42 @@ def test_call_from_dict(self, code):
 
 class TestLocalPythonExecutorSecurity:
     @pytest.mark.parametrize(
-        "additional_authorized_imports, expectation",
-        [([], pytest.raises(InterpreterError)), (["os"], does_not_raise())],
+        "additional_authorized_imports, expected_error",
+        [([], InterpreterError("Import of os is not allowed")), (["os"], None)],
     )
-    def test_vulnerability_import(self, additional_authorized_imports, expectation):
+    def test_vulnerability_import(self, additional_authorized_imports, expected_error):
         executor = LocalPythonExecutor(additional_authorized_imports)
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor("import os")
 
     @pytest.mark.parametrize(
-        "additional_authorized_imports, expectation",
-        [([], pytest.raises(InterpreterError)), (["builtins"], does_not_raise())],
+        "additional_authorized_imports, expected_error",
+        [([], InterpreterError("Import of builtins is not allowed")), (["builtins"], None)],
     )
-    def test_vulnerability_builtins(self, additional_authorized_imports, expectation):
+    def test_vulnerability_builtins(self, additional_authorized_imports, expected_error):
         executor = LocalPythonExecutor(additional_authorized_imports)
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor("import builtins")
 
     @pytest.mark.parametrize(
-        "additional_authorized_imports, expectation",
-        [([], pytest.raises(InterpreterError)), (["os", "sys"], does_not_raise())],
+        "additional_authorized_imports, expected_error",
+        [([], InterpreterError("Import of sys is not allowed")), (["os", "sys"], None)],
     )
-    def test_vulnerability_via_sys(self, additional_authorized_imports, expectation):
+    def test_vulnerability_via_sys(self, additional_authorized_imports, expected_error):
         executor = LocalPythonExecutor(additional_authorized_imports)
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor(
                 dedent(
                     """
@@ -1455,12 +1467,16 @@ def test_vulnerability_via_sys(self, additional_authorized_imports, expectation)
             )
 
     @pytest.mark.parametrize(
-        "additional_authorized_imports, expectation",
-        [(["importlib"], pytest.raises(InterpreterError)), (["importlib", "os"], does_not_raise())],
+        "additional_authorized_imports, expected_error",
+        [(["importlib"], InterpreterError("Forbidden return value: os")), (["importlib", "os"], None)],
     )
-    def test_vulnerability_via_importlib(self, additional_authorized_imports, expectation):
+    def test_vulnerability_via_importlib(self, additional_authorized_imports, expected_error):
         executor = LocalPythonExecutor(additional_authorized_imports)
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor(
                 dedent(
                     """
@@ -1472,11 +1488,11 @@ def test_vulnerability_via_importlib(self, additional_authorized_imports, expect
 
     @pytest.mark.parametrize("patch_builtin_import_module", [False, True])  # builtins_import.__module__ = None
     @pytest.mark.parametrize(
-        "additional_authorized_imports, additional_tools, expectation",
-        [([], [], pytest.raises(InterpreterError)), (["builtins", "os"], ["__import__"], does_not_raise())],
+        "additional_authorized_imports, additional_tools, expected_error",
+        [([], [], InterpreterError("Forbidden return value: builtins")), (["builtins", "os"], ["__import__"], None)],
     )
     def test_vulnerability_builtins_via_traceback(
-        self, patch_builtin_import_module, additional_authorized_imports, additional_tools, expectation, monkeypatch
+        self, patch_builtin_import_module, additional_authorized_imports, additional_tools, expected_error, monkeypatch
     ):
         if patch_builtin_import_module:
             monkeypatch.setattr("builtins.__import__.__module__", None)  # inspect.getmodule(func) = None
@@ -1485,7 +1501,11 @@ def test_vulnerability_builtins_via_traceback(
             from builtins import __import__
 
             executor.send_tools({"__import__": __import__})
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor(
                 dedent(
                     """
@@ -1502,11 +1522,11 @@ def test_vulnerability_builtins_via_traceback(
 
     @pytest.mark.parametrize("patch_builtin_import_module", [False, True])  # builtins_import.__module__ = None
     @pytest.mark.parametrize(
-        "additional_authorized_imports, additional_tools, expectation",
-        [([], [], pytest.raises(InterpreterError)), (["builtins", "os"], ["__import__"], does_not_raise())],
+        "additional_authorized_imports, additional_tools, expected_error",
+        [([], [], InterpreterError("Forbidden return value: builtins")), (["builtins", "os"], ["__import__"], None)],
     )
     def test_vulnerability_builtins_via_class_catch_warnings(
-        self, patch_builtin_import_module, additional_authorized_imports, additional_tools, expectation, monkeypatch
+        self, patch_builtin_import_module, additional_authorized_imports, additional_tools, expected_error, monkeypatch
     ):
         if patch_builtin_import_module:
             monkeypatch.setattr("builtins.__import__.__module__", None)  # inspect.getmodule(func) = None
@@ -1515,7 +1535,11 @@ def test_vulnerability_builtins_via_class_catch_warnings(
             from builtins import __import__
 
             executor.send_tools({"__import__": __import__})
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor(
                 dedent(
                     """
@@ -1533,12 +1557,16 @@ def test_vulnerability_builtins_via_class_catch_warnings(
 
     @pytest.mark.filterwarnings("ignore::DeprecationWarning")
     @pytest.mark.parametrize(
-        "additional_authorized_imports, expectation",
-        [([], pytest.raises(InterpreterError)), (["os"], does_not_raise())],
+        "additional_authorized_imports, expected_error",
+        [([], InterpreterError("Forbidden return value: os")), (["os"], None)],
     )
-    def test_vulnerability_load_module_via_builtin_importer(self, additional_authorized_imports, expectation):
+    def test_vulnerability_load_module_via_builtin_importer(self, additional_authorized_imports, expected_error):
         executor = LocalPythonExecutor(additional_authorized_imports)
-        with expectation:
+        with (
+            pytest.raises(type(expected_error), match=f".*{expected_error}")
+            if isinstance(expected_error, Exception)
+            else does_not_raise()
+        ):
             executor(
                 dedent(
                     """