Add @all-autogates test variant to wd_test

This adds support for running wd_test tests with all autogates enabled,
similar to the existing support in the internal repo.

Author: danlapid

build/kj_test.bzl

@@ -49,15 +49,15 @@ def kj_test(
     )
 
     sh_test(
-        name = test_name,
+        name = test_name + "@",
         size = size,
         srcs = ["//build/fixtures:kj_test.sh"],
         data = [cross_alias] + data,
         args = ["$(location " + cross_alias + ")"],
         tags = tags,
     )
     sh_test(
-        name = test_name + "@all-autogates-enabled",
+        name = test_name + "@all-autogates",
         size = size,
         env = {"WORKERD_ALL_AUTOGATES": "1"},
         srcs = ["//build/fixtures:kj_test.sh"],

build/wd_test.bzl

@@ -7,6 +7,8 @@ def wd_test(
         args = [],
         ts_deps = [],
         python_snapshot_test = False,
+        generate_default_variant = True,
+        generate_all_autogates_variant = True,
         **kwargs):
     """Rule to define tests that run `workerd test` with a particular config.
 
@@ -17,6 +19,8 @@ def wd_test(
      data: Additional files which the .capnp config file may embed. All TypeScript files will be compiled,
      their resulting files will be passed to the test as well. Usually TypeScript or Javascript source files.
      args: Additional arguments to pass to `workerd`. Typically used to pass `--experimental`.
+     autogates_variant: If True (default), generate an @all-autogates variant of the test.
+     default_variant: If True (default), generate the default (non-autogates) variant of the test.
     """
 
     # Add workerd binary to "data" dependencies.
@@ -57,14 +61,25 @@ def wd_test(
         "$(location {})".format(src),
     ] + args
 
-    _wd_test(
-        src = src,
-        name = name,
-        data = data,
-        args = args,
-        python_snapshot_test = python_snapshot_test,
-        **kwargs
-    )
+    if generate_default_variant:
+        _wd_test(
+            src = src,
+            name = name + "@",
+            data = data,
+            args = args,
+            python_snapshot_test = python_snapshot_test,
+            **kwargs
+        )
+
+    if generate_all_autogates_variant:
+        _wd_test(
+            src = src,
+            name = name + "@all-autogates",
+            data = data,
+            args = args + ["--all-autogates"],
+            python_snapshot_test = python_snapshot_test,
+            **kwargs
+        )
 
 WINDOWS_TEMPLATE = """
 @echo off
@@ -192,7 +207,7 @@ 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 = {}
+    environment = dict(ctx.attr.env)
     if ctx.attr.python_snapshot_test:
         environment["PYTHON_SAVE_SNAPSHOT_ARGS"] = ""
     if ctx.attr.load_snapshot:
@@ -272,6 +287,8 @@ _wd_test = rule(
         ),
         "python_snapshot_test": attr.bool(),
         "load_snapshot": attr.label(allow_single_file = True),
+        # Environment variables to set when running the test
+        "env": attr.string_dict(),
         # 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/node/tests/BUILD.bazel

@@ -335,6 +335,8 @@ wd_test(
     src = "dns-nodejs-test.wd-test",
     args = ["--experimental"],
     data = ["dns-nodejs-test.js"],
+    # TODO(jsg-rs): This test breaks when RUST_BACKED_NODE_DNS is enabled.
+    generate_all_autogates_variant = False,
     tags = ["requires-network"],
 )
 

src/workerd/api/tests/BUILD.bazel

@@ -341,6 +341,22 @@ wd_test(
     data = ["unsafe-test.js"],
 )
 
+# Test that verifies autogates are ENABLED - only runs in @all-autogates variant
+wd_test(
+    src = "autogate-enabled-test.wd-test",
+    args = ["--experimental"],
+    data = ["autogate-enabled-test.js"],
+    generate_default_variant = False,
+)
+
+# Test that verifies autogates are DISABLED - only runs in default variant
+wd_test(
+    src = "autogate-disabled-test.wd-test",
+    args = ["--experimental"],
+    data = ["autogate-disabled-test.js"],
+    generate_all_autogates_variant = False,
+)
+
 wd_test(
     src = "url-test.wd-test",
     args = ["--experimental"],

src/workerd/api/tests/autogate-disabled-test.js

@@ -0,0 +1,18 @@
+// Test that verifies autogates are DISABLED.
+// This test only runs in the default variant (all-autogates variant is disabled).
+
+import { strictEqual } from 'node:assert';
+import unsafe from 'workerd:unsafe';
+
+export const autogateDisabledTest = {
+  test() {
+    // In the default variant, WORKERD_ALL_AUTOGATES env var is NOT set,
+    // so isTestAutogateEnabled() should return false.
+    const isEnabled = unsafe.isTestAutogateEnabled();
+    strictEqual(
+      isEnabled,
+      false,
+      'TEST_WORKERD autogate should be DISABLED in default variant'
+    );
+  },
+};

src/workerd/api/tests/autogate-disabled-test.wd-test

@@ -0,0 +1,15 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "autogate-disabled-test",
+      worker = (
+        modules = [
+          (name = "worker", esModule = embed "autogate-disabled-test.js")
+        ],
+        compatibilityDate = "2024-01-01",
+        compatibilityFlags = ["nodejs_compat", "unsafe_module"],
+      )
+    ),
+  ],
+);

src/workerd/api/tests/autogate-enabled-test.js

@@ -0,0 +1,18 @@
+// Test that verifies autogates are ENABLED.
+// This test only runs in the @all-autogates variant (default variant is disabled).
+
+import { strictEqual } from 'node:assert';
+import unsafe from 'workerd:unsafe';
+
+export const autogateEnabledTest = {
+  test() {
+    // In the @all-autogates variant, WORKERD_ALL_AUTOGATES env var is set,
+    // so isTestAutogateEnabled() should return true.
+    const isEnabled = unsafe.isTestAutogateEnabled();
+    strictEqual(
+      isEnabled,
+      true,
+      'TEST_WORKERD autogate should be ENABLED in @all-autogates variant'
+    );
+  },
+};

src/workerd/api/tests/autogate-enabled-test.wd-test

@@ -0,0 +1,15 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "autogate-enabled-test",
+      worker = (
+        modules = [
+          (name = "worker", esModule = embed "autogate-enabled-test.js")
+        ],
+        compatibilityDate = "2024-01-01",
+        compatibilityFlags = ["nodejs_compat", "unsafe_module"],
+      )
+    ),
+  ],
+);

src/workerd/api/unsafe.c++

@@ -3,6 +3,7 @@
 #include <workerd/io/io-context.h>
 #include <workerd/jsg/jsg.h>
 #include <workerd/jsg/script.h>
+#include <workerd/util/autogate.h>
 
 namespace workerd::api {
 
@@ -203,6 +204,10 @@ jsg::Promise<void> UnsafeModule::abortAllDurableObjects(jsg::Lock& js) {
   return js.resolvedPromise();
 }
 
+bool UnsafeModule::isTestAutogateEnabled() {
+  return util::Autogate::isEnabled(util::AutogateKey::TEST_WORKERD);
+}
+
 #ifdef WORKERD_FUZZILLI
 void Fuzzilli::fuzzilli(jsg::Lock& js, jsg::Arguments<jsg::Value> args) {
   // Delegate to the fuzzilli_handler in fuzzilli.c++

src/workerd/api/unsafe.h

@@ -100,8 +100,13 @@ class UnsafeModule: public jsg::Object {
   UnsafeModule(jsg::Lock&, const jsg::Url&) {}
   jsg::Promise<void> abortAllDurableObjects(jsg::Lock& js);
 
+  // Returns true if the TEST_WORKERD autogate is enabled.
+  // This is used to verify that the all-autogates test variant is working correctly.
+  bool isTestAutogateEnabled();
+
   JSG_RESOURCE_TYPE(UnsafeModule) {
     JSG_METHOD(abortAllDurableObjects);
+    JSG_METHOD(isTestAutogateEnabled);
   }
 };