Move snapshot tests into workerd (#5205)

Gives us test coverage for memory snapshots in workerd

Author: hoodmane

build/wd_test.bzl

@@ -113,20 +113,20 @@ SH_RUNTEST_NORMAL = """"$@" -dTEST_TMPDIR=$TEST_TMPDIR"""
 # invoke workerd twice for snapshot tests.
 
 WINDOWS_RUNTEST_SNAPSHOT = """
-powershell -Command \"%* --python-save-snapshot\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
+powershell -Command \"%* --python-save-snapshot $ENV:PYTHON_SAVE_SNAPSHOT_ARGS\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
 set TEST_EXIT=!ERRORLEVEL!
 if !TEST_EXIT! EQU 0 (
-    powershell -Command \"%* --python-load-snapshot\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
+    powershell -Command \"%* --python-load-snapshot snapshot.bin\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
     set TEST_EXIT=!ERRORLEVEL!
 )
 """
 
 SH_RUNTEST_SNAPSHOT = """
 echo Creating Python Snapshot
-"$@" -dTEST_TMPDIR=$TEST_TMPDIR --python-save-snapshot
+"$@" -dTEST_TMPDIR=$TEST_TMPDIR --python-save-snapshot $PYTHON_SAVE_SNAPSHOT_ARGS
 echo ""
 echo Using Python Snapshot
-"$@" -dTEST_TMPDIR=$TEST_TMPDIR --python-load-snapshot
+"$@" -dTEST_TMPDIR=$TEST_TMPDIR --python-load-snapshot snapshot.bin
 """
 
 def _wd_test_impl(ctx):
@@ -191,12 +191,23 @@ def _wd_test_impl(ctx):
         # Include all file types that might contain testable code
         extensions = ["cc", "c++", "cpp", "cxx", "c", "h", "hh", "hpp", "hxx", "inc", "js", "ts", "mjs", "wd-test", "capnp"],
     )
+    environment = {}
+    if ctx.attr.python_snapshot_test:
+        environment["PYTHON_SAVE_SNAPSHOT_ARGS"] = ""
+    if ctx.attr.load_snapshot:
+        if ctx.attr.python_snapshot_test:
+            environment["PYTHON_SAVE_SNAPSHOT_ARGS"] = "--python-load-snapshot load_snapshot.bin"
+        f = ctx.attr.load_snapshot.files.to_list()[0]
+        runfiles = runfiles.merge(ctx.runfiles(symlinks = {"load_snapshot.bin": f}))
 
     return [
         DefaultInfo(
             executable = executable,
             runfiles = runfiles,
         ),
+        RunEnvironmentInfo(
+            environment = environment,
+        ),
         instrumented_files_info,
     ]
 
@@ -259,6 +270,7 @@ _wd_test = rule(
             default = "//src/workerd/api/node/tests:sidecar-supervisor",
         ),
         "python_snapshot_test": attr.bool(),
+        "load_snapshot": attr.label(allow_single_file = True),
         # A reference to the Windows platform label, needed for the implementation of wd_test
         "_platforms_os_windows": attr.label(default = "@platforms//os:windows"),
     },

src/workerd/api/pyodide/pyodide.h

@@ -61,7 +61,7 @@ struct PythonConfig {
   const PyodidePackageManager pyodidePackageManager;
   bool createSnapshot;
   bool createBaselineSnapshot;
-  bool loadSnapshotFromDisk;
+  kj::Maybe<kj::String> loadSnapshotFromDisk;
 };
 
 // A function to read a segment of the tar file into a buffer

src/workerd/server/server.h

@@ -78,8 +78,8 @@ class Server final: private kj::TaskSet::ErrorHandler {
   void setPythonCreateBaselineSnapshot() {
     pythonConfig.createBaselineSnapshot = true;
   }
-  void setPythonLoadSnapshot() {
-    pythonConfig.loadSnapshotFromDisk = true;
+  void setPythonLoadSnapshot(kj::String snapshot) {
+    pythonConfig.loadSnapshotFromDisk = kj::mv(snapshot);
   }
 
   // Runs the server using the given config.
@@ -126,7 +126,7 @@ class Server final: private kj::TaskSet::ErrorHandler {
     .pyodideDiskCacheRoot = kj::none,
     .createSnapshot = false,
     .createBaselineSnapshot = false,
-    .loadSnapshotFromDisk = false};
+    .loadSnapshotFromDisk = kj::none};
 
   bool experimental = false;
 

src/workerd/server/tests/python/BUILD.bazel

@@ -1,17 +1,7 @@
-load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
 load("//src/workerd/server/tests/python:import_tests.bzl", "gen_import_tests", "gen_rust_import_tests")
-load("//src/workerd/server/tests/python:py_wd_test.bzl", "py_wd_test")
-
-# pyodide_dev.capnp.bin represents a custom pyodide version "dev" that is generated
-# at build time using the latest contents of the src/pyodide directory.
-# This is used to run tests to ensure that they are always run against the latest build of
-# the Pyodide bundle.
-copy_file(
-    name = "pyodide_dev.capnp.bin@rule",
-    src = "//src/pyodide:pyodide.capnp.bin_cross",
-    out = "pyodide-bundle-cache/pyodide_dev.capnp.bin",
-    visibility = ["//visibility:public"],
-)
+load("//src/workerd/server/tests/python:py_wd_test.bzl", "py_wd_test", "python_test_setup")
+
+python_test_setup()
 
 py_wd_test("hello")
 
@@ -70,3 +60,15 @@ py_wd_test("vendor_dir_compat_flag")
 py_wd_test("multiprocessing")
 
 py_wd_test("default-class-with-legacy-global-handlers")
+
+py_wd_test(
+    "fastapi",
+    make_snapshot = False,
+    use_snapshot = "fastapi",
+)
+
+py_wd_test(
+    "numpy",
+    make_snapshot = False,
+    use_snapshot = "numpy",
+)

src/workerd/server/tests/python/fastapi/fastapi.wd-test

@@ -0,0 +1,16 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "fastapi",
+      worker = (
+        modules = [
+          (name = "worker.py", pythonModule = embed "worker.py"),
+          (name = "fastapi", pythonRequirement = "fastapi")
+        ],
+        compatibilityDate = "2025-08-14",
+        compatibilityFlags = [%PYTHON_FEATURE_FLAGS],
+      )
+    ),
+  ],
+);

src/workerd/server/tests/python/fastapi/worker.py

@@ -0,0 +1,7 @@
+import fastapi
+from workers import WorkerEntrypoint
+
+
+class Default(WorkerEntrypoint):
+    def test(self):
+        assert fastapi.__version__ in {"0.110.0", "0.116.1"}

src/workerd/server/tests/python/import_tests.bzl

@@ -45,6 +45,7 @@ def _test(name, directory, wd_test, py_file, python_version, **kwds):
         python_flags = [python_version],
         make_snapshot = False,
         args = ["--experimental", "--pyodide-package-disk-cache-dir", "../all_pyodide_wheels_%s" % pkg_tag],
+        skip_default_data = True,
         data = [py_file, "@all_pyodide_wheels_%s//:whls" % pkg_tag],
         **kwds
     )

src/workerd/server/tests/python/numpy/numpy.wd-test

@@ -0,0 +1,16 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "numpy",
+      worker = (
+        modules = [
+          (name = "worker.py", pythonModule = embed "worker.py"),
+          (name = "numpy", pythonRequirement = "numpy")
+        ],
+        compatibilityDate = "2025-08-14",
+        compatibilityFlags = [%PYTHON_FEATURE_FLAGS],
+      )
+    ),
+  ],
+);

src/workerd/server/tests/python/numpy/worker.py

@@ -0,0 +1,8 @@
+import numpy as np
+from workers import WorkerEntrypoint
+
+
+class Default(WorkerEntrypoint):
+    def test(self):
+        res = np.arange(12).reshape((3, -1))[::-2, ::-2]
+        assert str(res) == "[[11  9]\n [ 3  1]]"

src/workerd/server/tests/python/py_wd_test.bzl

@@ -1,3 +1,4 @@
+load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
 load("//:build/python_metadata.bzl", "BUNDLE_VERSION_INFO")
 load("//:build/wd_test.bzl", "wd_test")
@@ -12,16 +13,29 @@ def _py_wd_test_helper(
         name,
         src,
         python_flag,
-        make_snapshot,
         *,
-        args = [],
+        make_snapshot,
+        use_snapshot,
+        args,
+        data = [],
         **kwargs):
     name_flag = name + "_" + python_flag
     templated_src = name_flag.replace("/", "-") + "@template"
     templated_src = "/".join(src.split("/")[:-1] + [templated_src])
     flags = _get_enable_flags(python_flag)
     feature_flags_txt = ",".join(['"{}"'.format(flag) for flag in flags])
 
+    load_snapshot = None
+    if use_snapshot:
+        version_info = BUNDLE_VERSION_INFO[python_flag]
+
+        snapshot = version_info[use_snapshot + "_snapshot"]
+        data = data + [":python_snapshots"]
+        load_snapshot = snapshot
+
+    if load_snapshot and not make_snapshot:
+        args += ["--python-load-snapshot", "load_snapshot.bin"]
+
     expand_template(
         name = name_flag + "@rule",
         out = templated_src,
@@ -34,9 +48,56 @@ def _py_wd_test_helper(
         name = name_flag + "@",
         args = args,
         python_snapshot_test = make_snapshot,
+        data = data,
+        load_snapshot = load_snapshot,
         **kwargs
     )
 
+def _snapshot_file(snapshot):
+    if not snapshot:
+        return []
+    copy_file(
+        name = "pyodide-snapshot-%s@copy" % snapshot,
+        src = "@pyodide-snapshot-%s//file" % snapshot,
+        out = snapshot,
+        visibility = ["//visibility:public"],
+    )
+    return [":" + snapshot]
+
+def _snapshot_files(
+        baseline_snapshot = None,
+        numpy_snapshot = None,
+        fastapi_snapshot = None,
+        **_kwds):
+    result = []
+    result += _snapshot_file(baseline_snapshot)
+    result += _snapshot_file(numpy_snapshot)
+    result += _snapshot_file(fastapi_snapshot)
+    return result
+
+def python_test_setup():
+    # pyodide_dev.capnp.bin represents a custom pyodide version "dev" that is generated
+    # at build time using the latest contents of the src/pyodide directory.
+    # This is used to run tests to ensure that they are always run against the latest build of
+    # the Pyodide bundle.
+    copy_file(
+        name = "pyodide_dev.capnp.bin@rule",
+        src = "//src/pyodide:pyodide.capnp.bin_cross",
+        out = "pyodide-bundle-cache/pyodide_dev.capnp.bin",
+        visibility = ["//visibility:public"],
+    )
+    data = []
+    for x in BUNDLE_VERSION_INFO.values():
+        if x["name"] == "development":
+            continue
+        data += _snapshot_files(**x)
+
+    native.filegroup(
+        name = "python_snapshots",
+        data = data,
+        visibility = ["//visibility:public"],
+    )
+
 def py_wd_test(
         directory = None,
         *,
@@ -49,12 +110,16 @@ def py_wd_test(
         size = "enormous",
         tags = [],
         make_snapshot = True,
+        use_snapshot = None,
+        skip_default_data = False,
         **kwargs):
     if python_flags == "all":
         python_flags = BUNDLE_VERSION_INFO.keys()
     python_flags = [flag for flag in python_flags if flag not in skip_python_flags and flag in BUNDLE_VERSION_INFO]
-    if data == None and directory != None:
-        data = native.glob(
+    if data == None:
+        data = []
+    if directory and not skip_default_data:
+        data += native.glob(
             [
                 directory + "/**",
             ],
@@ -82,8 +147,10 @@ def py_wd_test(
             src,
             python_flag,
             make_snapshot = make_snapshot,
+            use_snapshot = use_snapshot,
             data = data,
             args = args,
             size = size,
             tags = tags,
+            **kwargs
         )