Add engine_error() context manager for testing plugin exceptions

src/python/pants/backend/awslambda/python/target_types_test.py

@@ -19,9 +19,8 @@
 from pants.backend.python.target_types import PythonLibrary, PythonRequirementLibrary
 from pants.backend.python.target_types_rules import rules as python_target_types_rules
 from pants.build_graph.address import Address
-from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.target import InjectedDependencies, InvalidFieldException
-from pants.testutil.rule_runner import QueryRule, RuleRunner
+from pants.testutil.rule_runner import EngineErrorTrap, QueryRule, RuleRunner
 
 
 @pytest.fixture
@@ -91,11 +90,13 @@ def assert_resolved(handler: str, *, expected: str, is_file: bool) -> None:
     assert_resolved("path.to.lambda:func", expected="path.to.lambda:func", is_file=False)
     assert_resolved("lambda.py:func", expected="project.lambda:func", is_file=True)
 
-    with pytest.raises(ExecutionError):
+    with EngineErrorTrap() as trap:
         assert_resolved("doesnt_exist.py:func", expected="doesnt matter", is_file=True)
+    assert "Unmatched glob" in str(trap.e)
     # Resolving >1 file is an error.
-    with pytest.raises(ExecutionError):
+    with EngineErrorTrap() as trap:
         assert_resolved("*.py:func", expected="doesnt matter", is_file=True)
+    assert isinstance(trap.e, InvalidFieldException)
 
 
 def test_inject_handler_dependency(rule_runner: RuleRunner, caplog) -> None:

src/python/pants/backend/go/util_rules/external_module_test.py

@@ -22,10 +22,9 @@
 from pants.backend.go.util_rules.go_pkg import ResolvedGoPackage
 from pants.engine.addresses import Address
 from pants.engine.fs import Digest, PathGlobs, Snapshot
-from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.process import ProcessExecutionFailure
 from pants.engine.rules import QueryRule
-from pants.testutil.rule_runner import RuleRunner
+from pants.testutil.rule_runner import EngineErrorTrap, RuleRunner
 
 
 @pytest.fixture
@@ -127,16 +126,13 @@ def assert_module(module: str, version: str, sample_file: str) -> None:
 
 def test_download_modules_missing_module(rule_runner: RuleRunner) -> None:
     input_digest = rule_runner.make_snapshot({"go.mod": GO_MOD, "go.sum": GO_SUM}).digest
-    with pytest.raises(ExecutionError) as exc:
+    with EngineErrorTrap() as trap:
         rule_runner.request(
             DownloadedModule,
             [DownloadedModuleRequest("some_project.org/project", "v1.1", input_digest)],
         )
-    underlying_exception = exc.value.wrapped_exceptions[0]
-    assert isinstance(underlying_exception, AssertionError)
-    assert "The module some_project.org/project@v1.1 was not downloaded" in str(
-        underlying_exception
-    )
+    assert isinstance(trap.e, AssertionError)
+    assert "The module some_project.org/project@v1.1 was not downloaded" in str(trap.e)
 
 
 def test_download_modules_invalid_go_sum(rule_runner: RuleRunner) -> None:
@@ -157,11 +153,10 @@ def test_download_modules_invalid_go_sum(rule_runner: RuleRunner) -> None:
             ),
         }
     ).digest
-    with pytest.raises(ExecutionError) as exc:
+    with EngineErrorTrap() as trap:
         rule_runner.request(AllDownloadedModules, [AllDownloadedModulesRequest(input_digest)])
-    underlying_exception = exc.value.wrapped_exceptions[0]
-    assert isinstance(underlying_exception, ProcessExecutionFailure)
-    assert "SECURITY ERROR" in str(underlying_exception)
+    assert isinstance(trap.e, ProcessExecutionFailure)
+    assert "SECURITY ERROR" in str(trap.e)
 
 
 def test_download_modules_missing_go_sum(rule_runner: RuleRunner) -> None:
@@ -183,10 +178,9 @@ def test_download_modules_missing_go_sum(rule_runner: RuleRunner) -> None:
             ),
         }
     ).digest
-    with pytest.raises(ExecutionError) as exc:
+    with EngineErrorTrap() as trap:
         rule_runner.request(AllDownloadedModules, [AllDownloadedModulesRequest(input_digest)])
-    underlying_exception = exc.value.wrapped_exceptions[0]
-    assert "`go.mod` and/or `go.sum` changed!" in str(underlying_exception)
+    assert "`go.mod` and/or `go.sum` changed!" in str(trap.e)
 
 
 def test_determine_external_package_info(rule_runner: RuleRunner) -> None:

src/python/pants/testutil/rule_runner.py

@@ -27,7 +27,7 @@
 from pants.engine.goal import Goal
 from pants.engine.internals import native_engine
 from pants.engine.internals.native_engine import PyExecutor
-from pants.engine.internals.scheduler import SchedulerSession
+from pants.engine.internals.scheduler import ExecutionError, SchedulerSession
 from pants.engine.internals.selectors import Get, Params
 from pants.engine.internals.session import SessionValues
 from pants.engine.process import InteractiveRunner
@@ -75,6 +75,53 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+class EngineErrorTrap:
+    """A context manager to catch `ExecutionError`s in tests and get the underlying exception.
+
+    Use like this:
+
+        with EngineErrorTrap() as trap:
+            rule_runner.request(OutputType, [input])
+        assert isinstance(trap.e, ValueError)
+
+    Will raise AssertionError if no ExecutionError occurred. Will re-raise all other exception
+    types.
+    """
+
+    _e: Exception | None = None
+
+    def __enter__(self) -> EngineErrorTrap:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is None:
+            raise AssertionError("Did not raise an exception! ExecutionError expected.")
+        if not issubclass(exc_type, ExecutionError):
+            # This is equivalent to re-raising the exception.
+            return None
+
+        if not len(exc_val.wrapped_exceptions) == 1:
+            formatted_errors = "\n\n".join(repr(e) for e in exc_val.wrapped_exceptions)
+            raise ValueError(
+                "Multiple underlying exceptions, but this helper function expected only one. "
+                "Use `with pytest.raises(ExecutionError)` directly and inspect "
+                "`exc.value.wrapped_exceptions`.\n\n"
+                f"Errors: {formatted_errors}"
+            )
+        self._e = exc_val.wrapped_exceptions[0]
+        return True
+
+    @property
+    def e(self) -> Exception:
+        if self._e is None:
+            raise AssertionError(
+                "`EngineErrorTrap.e` is not set. This happens when you try accessing the field "
+                "inside the context manager (i.e. the `with` block) and an ExecutionError has not "
+                "happened yet."
+            )
+        return self._e
+
+
 # -----------------------------------------------------------------------------------------------
 # `RuleRunner`
 # -----------------------------------------------------------------------------------------------