Make a JSON version of bazel_worker_test.

RELNOTES: None.
PiperOrigin-RevId: 332312350
diff --git a/src/test/java/com/google/devtools/build/lib/BUILD b/src/test/java/com/google/devtools/build/lib/BUILD
index 6b20cbe..2c071f1 100644
--- a/src/test/java/com/google/devtools/build/lib/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/BUILD
@@ -595,14 +595,18 @@
 
 java_library(
     name = "ExampleWorker-lib",
-    srcs = glob(["worker/ExampleWorker*.java"]),
+    srcs = glob(["worker/*ExampleWorker*.java"]),
     visibility = [
         "//src/test/shell/integration:__pkg__",
     ],
     deps = [
+        "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:worker_protocol_java_proto",
+        "//third_party:gson",
         "//third_party:guava",
+        "//third_party/protobuf:protobuf_java",
+        "//third_party/protobuf:protobuf_java_util",
     ],
 )
 
diff --git a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java
index f6a01be..7a31a62 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java
@@ -38,9 +38,7 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
-/**
- * An example implementation of a worker process that is used for integration tests.
- */
+/** An example implementation of a worker process that is used for integration tests. */
 public class ExampleWorker {
 
   static final Pattern FLAG_FILE_PATTERN = Pattern.compile("(?:@|--?flagfile=)(.+)");
@@ -67,7 +65,6 @@
       parser.parse(args);
       ExampleWorkerOptions workerOptions = parser.getOptions(ExampleWorkerOptions.class);
       Preconditions.checkState(workerOptions.persistentWorker);
-
       runPersistentWorker(workerOptions);
     } else {
       // This is a single invocation of the example that exits after it processed the request.
@@ -79,9 +76,18 @@
     PrintStream originalStdOut = System.out;
     PrintStream originalStdErr = System.err;
 
+    ExampleWorkerProtocol workerProtocol = null;
+    switch (workerOptions.workerProtocol) {
+      case JSON:
+        workerProtocol = new JsonExampleWorkerProtocolImpl(System.in, System.out);
+        break;
+      case PROTO:
+        workerProtocol = new ProtoExampleWorkerProtocolImpl(System.in, System.out);
+    }
+    Preconditions.checkNotNull(workerProtocol);
     while (true) {
       try {
-        WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);
+        WorkRequest request = workerProtocol.readRequest();
         if (request == null) {
           break;
         }
@@ -129,14 +135,12 @@
 
         if (poisoned) {
           baos.writeTo(System.out);
+          System.out.flush();
         } else {
-          WorkResponse.newBuilder()
-              .setOutput(baos.toString())
-              .setExitCode(exitCode)
-              .build()
-              .writeDelimitedTo(System.out);
+          WorkResponse response =
+              WorkResponse.newBuilder().setOutput(baos.toString()).setExitCode(exitCode).build();
+          workerProtocol.writeResponse(response);
         }
-        System.out.flush();
 
         if (workerOptions.exitAfter > 0 && workUnitCounter > workerOptions.exitAfter) {
           return;
diff --git a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java
index d541317..3e78481 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java
@@ -13,6 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.worker;
 
+import com.google.devtools.build.lib.actions.ExecutionRequirements;
+import com.google.devtools.common.options.EnumConverter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDocumentationCategory;
 import com.google.devtools.common.options.OptionEffectTag;
@@ -125,4 +127,21 @@
     help = "Instead of writing an error message to stdout, write it to stderr and terminate."
   )
   public boolean hardPoison;
+
+  /** Enum converter for --worker_protocol. */
+  public static class WorkerProtocolEnumConverter
+      extends EnumConverter<ExecutionRequirements.WorkerProtocolFormat> {
+    public WorkerProtocolEnumConverter() {
+      super(ExecutionRequirements.WorkerProtocolFormat.class, "worker protocol format option");
+    }
+  }
+
+  @Option(
+      name = "worker_protocol",
+      documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
+      effectTags = {OptionEffectTag.NO_OP},
+      defaultValue = "proto",
+      help = "The protocol (JSON or proto) to use for communication between this worker and Bazel.",
+      converter = WorkerProtocolEnumConverter.class)
+  public ExecutionRequirements.WorkerProtocolFormat workerProtocol;
 }
diff --git a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerProtocol.java b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerProtocol.java
new file mode 100644
index 0000000..34dd168
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerProtocol.java
@@ -0,0 +1,27 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.worker;
+
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
+import java.io.IOException;
+
+/** Example protocol for workers. */
+public interface ExampleWorkerProtocol extends AutoCloseable {
+  /** Reads a work request written to the worker's stdin. */
+  WorkRequest readRequest() throws IOException;
+
+  /** Writes a work response to the worker's stdout. */
+  void writeResponse(WorkResponse workResponse) throws IOException;
+}
diff --git a/src/test/java/com/google/devtools/build/lib/worker/JsonExampleWorkerProtocolImpl.java b/src/test/java/com/google/devtools/build/lib/worker/JsonExampleWorkerProtocolImpl.java
new file mode 100644
index 0000000..bde92cc
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/worker/JsonExampleWorkerProtocolImpl.java
@@ -0,0 +1,159 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.worker;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.devtools.build.lib.worker.WorkerProtocol.Input;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
+import com.google.gson.stream.JsonReader;
+import com.google.gson.stream.MalformedJsonException;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.util.JsonFormat;
+import com.google.protobuf.util.JsonFormat.Printer;
+import java.io.BufferedReader;
+import java.io.BufferedWriter;
+import java.io.EOFException;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
+import java.util.ArrayList;
+import java.util.List;
+
+/** Sample implementation of the Worker Protocol using JSON to communicate with Bazel. */
+public class JsonExampleWorkerProtocolImpl implements ExampleWorkerProtocol {
+  private final Printer jsonPrinter =
+      JsonFormat.printer().omittingInsignificantWhitespace().includingDefaultValueFields();
+  private final JsonReader reader;
+  private final BufferedWriter jsonWriter;
+
+  public JsonExampleWorkerProtocolImpl(InputStream stdin, OutputStream stdout) {
+    reader = new JsonReader(new BufferedReader(new InputStreamReader(stdin, UTF_8)));
+    reader.setLenient(true);
+    jsonWriter = new BufferedWriter(new OutputStreamWriter(stdout, UTF_8));
+  }
+
+  private static ArrayList<String> readArguments(JsonReader reader) throws IOException {
+    reader.beginArray();
+    ArrayList<String> arguments = new ArrayList<>();
+    while (reader.hasNext()) {
+      arguments.add(reader.nextString());
+    }
+    reader.endArray();
+    return arguments;
+  }
+
+  private static ArrayList<Input> readInputs(JsonReader reader) throws IOException {
+    reader.beginArray();
+    ArrayList<Input> inputs = new ArrayList<>();
+    while (reader.hasNext()) {
+      String digest = null;
+      String path = null;
+
+      reader.beginObject();
+      while (reader.hasNext()) {
+        String name = reader.nextName();
+        switch (name) {
+          case "digest":
+            if (digest != null) {
+              throw new IOException("Input cannot have more than one digest");
+            }
+            digest = reader.nextString();
+            break;
+          case "path":
+            if (path != null) {
+              throw new IOException("Input cannot have more than one path");
+            }
+            path = reader.nextString();
+            break;
+          default:
+            throw new IOException(name + " is an incorrect field in input");
+        }
+      }
+      reader.endObject();
+      inputs.add(
+          Input.newBuilder().setDigest(ByteString.copyFromUtf8(digest)).setPath(path).build());
+    }
+    reader.endArray();
+    return inputs;
+  }
+
+  @Override
+  public WorkRequest readRequest() throws IOException {
+    List<String> arguments = null;
+    List<Input> inputs = null;
+    Integer requestId = null;
+    try {
+      reader.beginObject();
+      while (reader.hasNext()) {
+        String name = reader.nextName();
+        switch (name) {
+          case "arguments":
+            if (arguments != null) {
+              throw new IOException("Work request cannot have more than one list of arguments");
+            }
+            arguments = readArguments(reader);
+            break;
+          case "inputs":
+            if (inputs != null) {
+              throw new IOException("Work request cannot have more than one list of inputs");
+            }
+            inputs = readInputs(reader);
+            break;
+          case "requestId":
+            if (requestId != null) {
+              throw new IOException("Work request cannot have more than one requestId");
+            }
+            requestId = reader.nextInt();
+            break;
+          default:
+            throw new IOException(name + " is an incorrect field in work request");
+        }
+      }
+      reader.endObject();
+    } catch (MalformedJsonException | IllegalStateException | EOFException e) {
+      throw new IOException(e);
+    }
+
+    WorkRequest.Builder requestBuilder = WorkRequest.newBuilder();
+    if (arguments != null) {
+      requestBuilder.addAllArguments(arguments);
+    }
+    if (inputs != null) {
+      requestBuilder.addAllInputs(inputs);
+    }
+    if (requestId != null) {
+      requestBuilder.setRequestId(requestId);
+    }
+    return requestBuilder.build();
+  }
+
+  @Override
+  public void writeResponse(WorkResponse response) throws IOException {
+    jsonPrinter.appendTo(response, jsonWriter);
+    jsonWriter.flush();
+  }
+
+  @Override
+  public void close() {
+    try {
+      jsonWriter.close();
+    } catch (IOException e) {
+      System.err.printf("Could not close json writer. %s", e);
+    }
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/worker/ProtoExampleWorkerProtocolImpl.java b/src/test/java/com/google/devtools/build/lib/worker/ProtoExampleWorkerProtocolImpl.java
new file mode 100644
index 0000000..e0ee1cb
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/worker/ProtoExampleWorkerProtocolImpl.java
@@ -0,0 +1,50 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.worker;
+
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
+import com.google.protobuf.InvalidProtocolBufferException;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+/** Sample implementation of the Worker Protocol using Proto to communiocate with Bazel. */
+public class ProtoExampleWorkerProtocolImpl implements ExampleWorkerProtocol {
+  private final InputStream stdin;
+  private final OutputStream stdout;
+
+  public ProtoExampleWorkerProtocolImpl(InputStream stdin, OutputStream stdout) {
+    this.stdin = stdin;
+    this.stdout = stdout;
+  }
+
+  @Override
+  public WorkRequest readRequest() throws IOException {
+    try {
+      return WorkRequest.parseDelimitedFrom(stdin);
+    } catch (InvalidProtocolBufferException e) {
+      throw new IOException(e);
+    }
+  }
+
+  @Override
+  public void writeResponse(WorkResponse response) throws IOException {
+    response.writeDelimitedTo(stdout);
+    stdout.flush();
+  }
+
+  @Override
+  public void close() {}
+}
diff --git a/src/test/shell/integration/BUILD b/src/test/shell/integration/BUILD
index 30a0887..12a6945 100644
--- a/src/test/shell/integration/BUILD
+++ b/src/test/shell/integration/BUILD
@@ -465,6 +465,7 @@
     args = [
         "--worker_sandboxing=no",
         "non-sandboxed",
+        "proto",
     ],
     data = [
         ":test-deps",
@@ -477,12 +478,31 @@
 )
 
 sh_test(
+    name = "bazel_json_worker_test",
+    size = "large",
+    srcs = ["bazel_worker_test.sh"],
+    args = [
+        "--worker_sandboxing=no",
+        "non-sandboxed",
+        "json",
+    ],
+    data = [
+        ":test-deps",
+        "//src/test/java/com/google/devtools/build/lib:ExampleWorker_deploy.jar",
+    ],
+    tags = [
+        "no_windows",
+    ],
+)
+
+sh_test(
     name = "bazel_worker_multiplexer_test",
     size = "large",
     srcs = ["bazel_worker_multiplexer_test.sh"],
     args = [
         "--worker_sandboxing=no",
         "non-sandboxed",
+        "proto",
     ],
     data = [
         ":test-deps",
@@ -501,6 +521,7 @@
     args = [
         "--worker_sandboxing",
         "sandboxed",
+        "proto",
     ],
     data = [
         ":test-deps",
diff --git a/src/test/shell/integration/bazel_worker_test.sh b/src/test/shell/integration/bazel_worker_test.sh
index 0553c18..9b4f11c 100755
--- a/src/test/shell/integration/bazel_worker_test.sh
+++ b/src/test/shell/integration/bazel_worker_test.sh
@@ -20,7 +20,8 @@
 set -u
 ADDITIONAL_BUILD_FLAGS=$1
 WORKER_TYPE_LOG_STRING=$2
-shift 2
+WORKER_PROTOCOL=$3
+shift 3
 
 # Load the test setup defined in the parent directory
 CURRENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -34,6 +35,7 @@
 
 add_to_bazelrc "build -s"
 add_to_bazelrc "build --spawn_strategy=worker,standalone"
+add_to_bazelrc "build --experimental_allow_json_worker_protocol"
 add_to_bazelrc "build --worker_verbose --worker_max_instances=1"
 add_to_bazelrc "build --debug_print_action_contexts"
 add_to_bazelrc "build ${ADDITIONAL_BUILD_FLAGS}"
@@ -100,7 +102,7 @@
   mkdir worker_data_dir
   echo "veryexample" > worker_data_dir/more_data.txt
 
-  cat >work.bzl <<'EOF'
+  cat >work.bzl <<EOF
 def _impl(ctx):
   worker = ctx.executable.worker
   output = ctx.outputs.out
@@ -126,7 +128,7 @@
     argfile_inputs.append(argfile)
     argfile_arguments.append("@" + argfile.path)
 
-  execution_requirements = {"supports-workers": "1"}
+  execution_requirements = {"supports-workers": "1", "requires-worker-protocol": "$WORKER_PROTOCOL"}
   if ctx.attr.worker_key_mnemonic:
     execution_requirements["worker-key-mnemonic"] = ctx.attr.worker_key_mnemonic
 
@@ -182,12 +184,14 @@
 work(
   name = "hello_world",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["hello world"],
 )
 
 work(
   name = "hello_world_uppercase",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--uppercase", "hello world"],
 )
 EOF
@@ -207,6 +211,7 @@
 work(
   name = "hello_world",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   action_mnemonic = "Hello",
   worker_key_mnemonic = "SharedWorker",
   args = ["--write_uuid"],
@@ -215,6 +220,7 @@
 work(
   name = "goodbye_world",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   action_mnemonic = "Goodbye",
   worker_key_mnemonic = "SharedWorker",
   args = ["--write_uuid"],
@@ -234,6 +240,7 @@
 work(
   name = "multi_hello_world",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["hello", "world", "nice", "to", "meet", "you"],
   multiflagfiles = True,
 )
@@ -246,10 +253,11 @@
 
 function test_workers_quit_after_build() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -268,11 +276,11 @@
 
 function test_build_succeeds_even_if_worker_exits() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
-  worker_args = ["--exit_after=1"],
+  worker_args = ["--exit_after=1", "--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -284,17 +292,20 @@
   bazel build --worker_verbose :hello_world_2 &> $TEST_log \
     || fail "build failed"
 
-  expect_log "Work worker (id 2) has unexpectedly died with exit code 0."
+  expect_log "Work worker (id [0-9]\+) has unexpectedly died with exit code 0."
 }
 
 function test_build_fails_if_worker_dies_during_action() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
-  worker_args = ["--exit_during=1"],
-  args = ["--write_uuid", "--write_counter"],
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}","--exit_during=1"],
+  args = [
+    "--write_uuid",
+    "--write_counter",
+  ],
 ) for idx in range(10)]
 EOF
 
@@ -306,10 +317,11 @@
 
 function test_worker_restarts_when_worker_binary_changes() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -351,10 +363,11 @@
 
 function test_worker_restarts_when_worker_runfiles_change() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -394,11 +407,11 @@
 # protobuf, it must be killed and a helpful error message should be printed.
 function test_build_fails_when_worker_returns_junk() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
-  worker_args = ["--poison_after=1"],
+  worker_args = ["--poison_after=1", "--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -418,10 +431,11 @@
 
 function test_input_digests() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--print_inputs"],
   srcs = [":input.txt"],
 ) for idx in range(10)]
@@ -454,10 +468,11 @@
 
 function test_worker_verbose() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -471,10 +486,11 @@
 
 function test_logs_are_deleted_on_server_restart() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF
@@ -497,9 +513,38 @@
     || fail "Worker log was not deleted"
 }
 
+function test_requires_worker_protocol_missing_defaults_to_proto {
+  prepare_example_worker
+  cat >>BUILD <<EOF
+work(
+  name = "hello_world_proto",
+  worker = ":worker",
+  worker_args = ["--worker_protocol=proto"],
+  args = ["hello world"],
+)
+work(
+  name = "hello_world_json",
+  worker = ":worker",
+  worker_args = ["--worker_protocol=json"],
+)
+EOF
+
+  sed -i.bak 's/=execution_requirements/={"supports-workers": "1"}/g' work.bzl
+  rm -f work.bzl.bak
+
+  bazel build :hello_world_proto &> $TEST_log \
+    || fail "build failed"
+  assert_equals "hello world" "$(cat $BINS/hello_world_proto.out)"
+
+  bazel build :hello_world_json &> $TEST_log \
+    && fail "expected proto build with json worker to fail" || true
+}
+
 function test_missing_execution_requirements_fallback_to_standalone() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  # This test ignores the WORKER_PROTOCOL test arg since it doesn't use the
+  # persistent worker when execution falls back to standalone.
+  cat >>BUILD <<EOF
 work(
   name = "hello_world",
   worker = ":worker",
@@ -523,10 +568,11 @@
 
 function test_environment_is_clean() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 work(
   name = "hello_world",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["--print_env"],
 )
 EOF
@@ -546,6 +592,7 @@
 work(
   name = "hello_clean",
   worker = ":worker",
+  worker_args = ["--worker_protocol=${WORKER_PROTOCOL}"],
   args = ["hello clean"],
 )
 EOF
@@ -563,11 +610,15 @@
 
 function test_crashed_worker_causes_log_dump() {
   prepare_example_worker
-  cat >>BUILD <<'EOF'
+  cat >>BUILD <<EOF
 [work(
   name = "hello_world_%s" % idx,
   worker = ":worker",
-  worker_args = ["--poison_after=1", "--hard_poison"],
+  worker_args = [
+    "--poison_after=1",
+    "--hard_poison",
+    "--worker_protocol=${WORKER_PROTOCOL}"
+  ],
   args = ["--write_uuid", "--write_counter"],
 ) for idx in range(10)]
 EOF