chore: Show correct test failure error context

packages/django-cf/tests/conftest.py

@@ -198,4 +198,6 @@ def dev_server(
 
 
 def register_in_worker_suites(namespace: dict, src_dir: Path) -> None:
-    register_testlib_suites(namespace, src_dir)
+    register_testlib_suites(
+        namespace, src_dir, source_roots=[PACKAGE_DIR, WORKERS_RUNTIME_SDK]
+    )

packages/runtime-sdk/tests/conftest.py

@@ -113,4 +113,10 @@ def register_in_worker_suites(
     src_dir: Path,
     marks: dict[str, pytest.MarkDecorator] | None = None,
 ) -> None:
-    register_testlib_suites(namespace, src_dir, marks=marks, class_name=str.upper)
+    register_testlib_suites(
+        namespace,
+        src_dir,
+        marks=marks,
+        class_name=str.upper,
+        source_roots=[WORKERS_RUNTIME_SDK],
+    )

packages/testlib/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "testlib"
 version = "0.0.0"
 requires-python = ">=3.11"
-dependencies = ["pytest"]
+dependencies = ["pytest", "pyodide-tblib>=3.2.3"]
 
 [project.optional-dependencies]
 host = ["requests"]

packages/testlib/testlib/entrypoint.py

@@ -12,6 +12,8 @@
 from pyodide.webloop import WebLoop
 from workers import Response, WorkerEntrypoint
 
+from .tracebacks import dump_exception
+
 
 class EnvPlugin:
     def __init__(self, env):
@@ -22,6 +24,18 @@ def env(self):
         return self._env
 
 
+def _describe_call_exception(item, call):
+    """Structured description of the exception recorded in ``call``.
+
+    Uses pytest's own traceback pruning so the host only sees the frames pytest
+    would have shown had the test run there directly.
+    """
+    excinfo = call.excinfo
+    traceback_filter = getattr(item, "_traceback_filter", excinfo.traceback.filter)
+    entries = [entry._rawentry for entry in traceback_filter(excinfo)]
+    return {"pickle": dump_exception(excinfo.value, entries), "when": call.when}
+
+
 class ResultCollector:
     """Record each pytest result under its host-suite test name."""
 
@@ -56,8 +70,10 @@ def pytest_runtest_makereport(self, item, call):
         elif report.failed:
             result = {"status": "failed", "traceback": report.longreprtext}
             self.results[key] = result
+            excinfo = call.excinfo
+            if excinfo is not None:
+                result["exception"] = _describe_call_exception(item, call)
             if report.when == "call":
-                excinfo = call.excinfo
                 if excinfo is None:
                     # e.g. a strict xfail that unexpectedly passed.
                     result["error"] = report.longreprtext or "unknown error"

packages/testlib/testlib/host.py

@@ -8,14 +8,16 @@
 import socket
 import subprocess
 import time
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Literal, TypedDict
+from typing import Any, Literal, NotRequired, TypedDict
 
 import pytest
 import requests
 
+from .tracebacks import WorkerException, load_exception
+
 SUITE_CONNECT_TIMEOUT = 10
 SUITE_READ_TIMEOUT = 300
 
@@ -67,6 +69,7 @@ class InWorkerTestResult(TypedDict):
     error: str
     traceback: str
     reason: str
+    exception: NotRequired[WorkerException]
 
 
 SuiteResults = dict[str, InWorkerTestResult]
@@ -187,8 +190,10 @@ def get_suite_results(server: str, suite: str) -> SuiteResults | str:
     return response.json()
 
 
-def _make_test(suite: str, test_name: str) -> Callable:
+def _make_test(suite: str, test_name: str, source_roots: tuple[Path, ...]) -> Callable:
     def test_fn(self: Any, dev_server: str) -> None:
+        # Hide this frame: the interesting traceback is the one from the worker.
+        __tracebackhide__ = True
         results = get_suite_results(dev_server, suite)
         if isinstance(results, str):
             pytest.fail(results)
@@ -199,10 +204,16 @@ def test_fn(self: Any, dev_server: str) -> None:
         )
         if result["status"] == "skipped":
             pytest.skip(result.get("reason", ""))
-        if result["status"] == "failed":
-            pytest.fail(result["error"])
-        if result["status"] == "error":
-            pytest.fail(f"{result['error']}\n{result.get('traceback', '')}")
+        if result["status"] not in ("failed", "error"):
+            return
+        exception = result.get("exception")
+        if exception is None:
+            pytest.fail(f"{result['error']}\n{result.get('traceback', '')}".rstrip())
+        exc = load_exception(exception, source_roots)
+        when = exception.get("when", "call")
+        if when != "call":
+            exc.add_note(f"raised in the worker during test {when}")
+        raise exc
 
     test_fn.__name__ = f"test_{test_name}"
     return test_fn
@@ -235,8 +246,16 @@ def register_in_worker_suites(
     *,
     marks: dict[str, pytest.MarkDecorator] | None = None,
     class_name: Callable[[str], str] | None = None,
+    source_roots: Sequence[Path] = (),
 ) -> None:
-    """Expose each in-worker test as an individual host-side pytest test."""
+    """Expose each in-worker test as an individual host-side pytest test.
+
+    ``source_roots`` lists extra host directories (besides ``src_dir``) that
+    hold copies of code running inside the worker, e.g. a package's source
+    tree that gets vendored into ``python_modules``. Traceback frames from
+    the worker are remapped onto them so pytest can show source lines.
+    """
+    roots = (src_dir, *source_roots)
     for module_path in sorted(src_dir.glob("test_*.py")):
         suite = module_path.stem[len("test_") :]
         generated_class_name = (
@@ -248,7 +267,7 @@ def register_in_worker_suites(
             f"Test{generated_class_name}",
             (),
             {
-                f"test_{name}": _make_test(suite, name)
+                f"test_{name}": _make_test(suite, name, roots)
                 for name in discover_test_names(module_path)
             },
         )

packages/testlib/testlib/tracebacks.py

@@ -0,0 +1,134 @@
+"""Ship exceptions raised inside the worker to the host with pickle.
+
+The worker pickles the exception with tblib's pickling support, which carries
+the traceback, the ``__cause__``/``__context__`` chain and ``__notes__``. Hooks
+on both ends keep that from failing on things pickle cannot handle.
+
+The host unpickles and raises the result, so pytest renders the failure exactly
+as if the test had run here.
+"""
+
+from __future__ import annotations
+
+import base64
+import functools
+import pickle
+from collections.abc import Sequence
+from io import BytesIO
+from pathlib import Path
+from types import TracebackType
+from typing import NotRequired, TypedDict
+
+from tblib import pickling_support
+
+
+class WorkerException(TypedDict):
+    """Pickled exception from the worker, base64-encoded."""
+
+    pickle: str
+    when: NotRequired[str]
+
+
+def _safe_repr(value: object) -> str:
+    try:
+        return repr(value)
+    except Exception:
+        return object.__repr__(value)
+
+
+# --- worker side: pickling ---------------------------------------------------
+
+
+class _Pickler(pickle.Pickler):
+    def reducer_override(self, obj):
+        if isinstance(obj, BaseException):
+            func, args, *state = pickling_support.pickle_exception(obj)
+            return (_unpickle_exception, (func, args, str(obj)), *state)
+        if isinstance(obj, TracebackType):
+            return pickling_support.pickle_traceback(obj)
+        if isinstance(obj, type):
+            try:
+                pickle.dumps(obj)
+            except Exception:
+                return (_exception_class, (obj.__qualname__, obj.__module__))
+            return NotImplemented
+        if type(obj).__module__ in ["builtins", "tblib"]:
+            return NotImplemented
+        return (str, (_safe_repr(obj),))
+
+
+def _link(entries: Sequence[TracebackType]) -> TracebackType | None:
+    """Chain non-contiguous traceback entries into one traceback."""
+    tb = None
+    for entry in reversed(entries):
+        tb = TracebackType(tb, entry.tb_frame, entry.tb_lasti, entry.tb_lineno)
+    return tb
+
+
+def dump_exception(exc: BaseException, entries: Sequence[TracebackType]) -> str:
+    """Pickle ``exc`` for the host; ``entries`` replaces its traceback if given."""
+    exc = exc.with_traceback(_link(entries))
+    buffer = BytesIO()
+    _Pickler(buffer, protocol=pickle.HIGHEST_PROTOCOL).dump(exc)
+    return base64.b64encode(buffer.getvalue()).decode()
+
+
+# --- host side: unpickling ---------------------------------------------------
+
+
+def _exception_class(qualname: str, module: str) -> type[Exception]:
+    """Dummy exception class. Pickle can't reference @functools.cache wrappers by name."""
+    return _make_exception_class(qualname, module)
+
+
+@functools.cache
+def _make_exception_class(qualname: str, module: str) -> type[Exception]:
+    return type(
+        qualname.rsplit(".", 1)[-1],
+        (Exception,),
+        {"__module__": module, "__qualname__": qualname, "_testlib_stand_in": True},
+    )
+
+
+def _unpickle_exception(func, args, message):
+    exc = func(*args)
+    if getattr(type(exc), "_testlib_stand_in", False):
+        # A stand-in lacks the original __str__, so make str(exc) the message
+        # the worker printed
+        exc.args = (message,)
+    return exc
+
+
+@functools.cache
+def _locate_source(filename: str, roots: tuple[Path, ...]) -> str:
+    """Map a worker-side path onto the host file it was copied from."""
+    if not filename.startswith("/session/metadata/"):
+        return filename
+    parts = Path(filename).parts[3:]
+    if parts[0] == "python_modules":
+        parts = parts[1:]
+    for root in roots:
+        candidate = root.joinpath(*parts)
+        if candidate.is_file():
+            return str(candidate)
+    return filename
+
+
+class _Unpickler(pickle.Unpickler):
+    def find_class(self, module: str, name: str):
+        try:
+            return super().find_class(module, name)
+        except (ImportError, AttributeError):
+            # The only globals the worker references that we may lack are
+            # exception classes; everything else was reduced to builtins.
+            return _exception_class(name, module)
+
+
+def load_exception(
+    payload: WorkerException, source_roots: Sequence[Path] = ()
+) -> BaseException:
+    """Recreate an exception pickled by ``dump_exception``."""
+    roots = tuple(source_roots)
+    data = base64.b64decode(payload["pickle"])
+    with pickling_support.map_filenames(lambda name: _locate_source(name, roots)):
+        return _Unpickler(BytesIO(data)).load()