diff --git a/rules_bazel/java/java_diff_test.bzl b/rules_bazel/java/java_diff_test.bzl
deleted file mode 100644
index 626c33b251..0000000000
--- a/rules_bazel/java/java_diff_test.bzl
+++ /dev/null
@@ -1,84 +0,0 @@
-def _overwrite_golden_impl(ctx):
-    test_class = ctx.attr.test_class
-    inputs = ctx.files.srcs
-    goldens_output_zip = ctx.outputs.goldens_output_zip
-    test_runner = ctx.executable.test_runner
-
-    # Generate the goldens from tests.
-    generate_goldens_script = """
-    mkdir local_tmp
-    TEST_OUTPUT_HOME="$(pwd)/local_tmp" {test_runner_path} $@
-    cd local_tmp
-    # Zip all files under local_tmp with all nested parent folders except for local_tmp itself.
-    # Zip files because there are cases that one Junit test can produce multiple goldens.
-    zip -r ../{goldens_output_zip} .
-    """.format(
-        test_runner_path = test_runner.path,
-        goldens_output_zip = goldens_output_zip.path,
-    )
-
-    ctx.actions.run_shell(
-        inputs = inputs,
-        outputs = [goldens_output_zip],
-        arguments = [test_class],
-        tools = [test_runner],
-        command = generate_goldens_script,
-    )
-
-    # Overwrite the goldens.
-    golden_update_script_content = """
-    #!/bin/bash
-    cd ${{BUILD_WORKSPACE_DIRECTORY}}
-    unzip -ao {goldens_output_zip} -d src/test/java
-    """.format(
-        goldens_output_zip = goldens_output_zip.path,
-    )
-    ctx.actions.write(
-        output = ctx.outputs.golden_update_script,
-        content = golden_update_script_content,
-        is_executable = True,
-    )
-    return [DefaultInfo(executable = ctx.outputs.golden_update_script)]
-
-overwrite_golden = rule(
-    attrs = {
-        "test_class": attr.string(mandatory = True),
-        "srcs": attr.label_list(
-            allow_files = True,
-            mandatory = True,
-        ),
-        "test_runner": attr.label(
-            mandatory = True,
-            executable = True,
-            cfg = "host",
-        ),
-    },
-    outputs = {
-        "goldens_output_zip": "%{name}.zip",
-        "golden_update_script": "%{name}.sh",
-    },
-    executable = True,
-    implementation = _overwrite_golden_impl,
-)
-
-def golden_update(
-        name,
-        srcs,
-        test_class,
-        data = [],
-        deps = []):
-    golden_junit_runner_name = "%s_junit_runner" % name
-    native.java_binary(
-        name = golden_junit_runner_name,
-        srcs = srcs,
-        data = data,
-        main_class = "com.google.api.generator.test.framework.SingleJUnitTestRunner",
-        deps = ["//src/test/java/com/google/api/generator/test/framework:junit_runner"] + deps,
-    )
-
-    overwrite_golden(
-        name = name,
-        test_class = test_class,
-        test_runner = ":%s" % golden_junit_runner_name,
-        srcs = srcs + data,
-    )
diff --git a/rules_bazel/unit_golden_update.bzl b/rules_bazel/unit_golden_update.bzl
new file mode 100644
index 0000000000..4d6ce13d46
--- /dev/null
+++ b/rules_bazel/unit_golden_update.bzl
@@ -0,0 +1,20 @@
+# Macro expanding to two rules:
+#   (1) java_library(): compiles given .java files
+#   (2) sh_binary(): script to update .golden (depends on (1))
+def golden_update(name, srcs, test_class, data, deps):
+    srcs_name = "%s_srcs" % name
+    native.java_library(
+        name = srcs_name,
+        srcs = srcs,
+        deps = deps,
+    )
+
+    native.sh_binary(
+        name = name,
+        srcs = ["//scripts:update_unit_golden.sh"],
+        args = [test_class],
+        data = data + deps + [
+            "//src/test/java/com/google/api/generator/test/framework:junit_runner",
+            srcs_name,
+        ],
+    )
diff --git a/scripts/BUILD.bazel b/scripts/BUILD.bazel
index 9f888ffad7..3d33717a25 100644
--- a/scripts/BUILD.bazel
+++ b/scripts/BUILD.bazel
@@ -1,3 +1,7 @@
 package(default_visibility = ["//test/integration:__pkg__"])
 
-exports_files(["diff_gen_and_golden.sh", "update_golden.sh"])
+exports_files([
+    "diff_gen_and_golden.sh",
+    "update_golden.sh",
+    "update_unit_golden.sh",
+])
diff --git a/scripts/update_unit_golden.sh b/scripts/update_unit_golden.sh
new file mode 100755
index 0000000000..7ac3b9a84c
--- /dev/null
+++ b/scripts/update_unit_golden.sh
@@ -0,0 +1,9 @@
+#!/bin/sh
+
+set -o errexit
+
+TEST_CLASS=$1
+CLASSPATH=$( find . -name '*.jar' ! -name 'liblite.jar' | xargs echo | tr ' ' ':' )
+
+TEST_OUTPUT_HOME=${BUILD_WORKSPACE_DIRECTORY}/src/test/java \
+  java -cp ${CLASSPATH} com.google.api.generator.test.framework.SingleJUnitTestRunner ${TEST_CLASS}
diff --git a/src/test/java/com/google/api/generator/engine/BUILD.bazel b/src/test/java/com/google/api/generator/engine/BUILD.bazel
index e52ef62387..42b2070517 100644
--- a/src/test/java/com/google/api/generator/engine/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/engine/BUILD.bazel
@@ -1,5 +1,5 @@
 load("@rules_java//java:defs.bzl", "java_test")
-load("//:rules_bazel/java/java_diff_test.bzl", "golden_update")
+load("//:rules_bazel/unit_golden_update.bzl", "golden_update")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -26,20 +26,18 @@ TEST_DEPS = [
 
 [java_test(
     name = test_name,
-    srcs = ["{0}.java".format(test_name)],
+    srcs = ["%s.java" % test_name],
     data = ["//src/test/java/com/google/api/generator/engine/goldens:goldens_files"],
-    test_class = "com.google.api.generator.engine.{0}".format(test_name),
+    test_class = "com.google.api.generator.engine.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
 
-TEST_CLASS_NAME = "com.google.api.generator.engine.JavaCodeGeneratorTest"
-
 # Run `bazel run src/test/java/com/google/api/generator/engine:JavaCodeGeneratorTest_update`
 # to update goldens as expected generated code.
 [golden_update(
-    name = "{0}_update".format(test_name),
-    srcs = ["{0}.java".format(test_name)],
+    name = "%s_update" % test_name,
+    srcs = ["%s.java" % test_name],
     data = ["//src/test/java/com/google/api/generator/engine/goldens:goldens_files"],
-    test_class = "com.google.api.generator.engine.{0}".format(test_name),
+    test_class = "com.google.api.generator.engine.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
diff --git a/src/test/java/com/google/api/generator/gapic/composer/BUILD.bazel b/src/test/java/com/google/api/generator/gapic/composer/BUILD.bazel
index 1ad263ea22..358c924514 100644
--- a/src/test/java/com/google/api/generator/gapic/composer/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/gapic/composer/BUILD.bazel
@@ -1,14 +1,12 @@
 load("@rules_java//java:defs.bzl", "java_proto_library", "java_test")
-load("//:rules_bazel/java/java_diff_test.bzl", "golden_update")
+load("//:rules_bazel/unit_golden_update.bzl", "golden_update")
 
 package(default_visibility = ["//visibility:public"])
 
-UPDATE_GOLDENS_TESTS = [
+TESTS = [
     "ComposerTest",
 ]
 
-TESTS = UPDATE_GOLDENS_TESTS
-
 TEST_DEPS = [
     ":common_resources_java_proto",
     "//:service_config_java_proto",
@@ -60,14 +58,14 @@ java_proto_library(
 [java_test(
     name = test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ],
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
 
@@ -76,15 +74,15 @@ java_proto_library(
 # `ServiceClient*` tests are not supported now since they are still in active development.
 
 [golden_update(
-    name = "{0}_update".format(test_name),
+    name = "%s_update" % test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ],
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.%s" % test_name,
     deps = TEST_DEPS,
-) for test_name in UPDATE_GOLDENS_TESTS]
+) for test_name in TESTS]
diff --git a/src/test/java/com/google/api/generator/gapic/composer/common/BUILD.bazel b/src/test/java/com/google/api/generator/gapic/composer/common/BUILD.bazel
index 28a4bfc8b4..1faa7bde6c 100644
--- a/src/test/java/com/google/api/generator/gapic/composer/common/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/gapic/composer/common/BUILD.bazel
@@ -1,5 +1,5 @@
 load("@rules_java//java:defs.bzl", "java_proto_library", "java_test")
-load("//:rules_bazel/java/java_diff_test.bzl", "golden_update")
+load("//:rules_bazel/unit_golden_update.bzl", "golden_update")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -70,14 +70,14 @@ java_proto_library(
 [java_test(
     name = test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ],
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/common/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.common.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.common.%s" % test_name,
     deps = TEST_DEPS + [":common"],
 ) for test_name in TESTS]
 
@@ -86,15 +86,15 @@ java_proto_library(
 # `ServiceClient*` tests are not supported now since they are still in active development.
 
 [golden_update(
-    name = "{0}_update".format(test_name),
+    name = "%s_update" % test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ],
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/common/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.common.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.common.%s" % test_name,
     deps = TEST_DEPS + [":common"],
 ) for test_name in UPDATE_GOLDENS_TESTS]
diff --git a/src/test/java/com/google/api/generator/gapic/composer/grpc/BUILD.bazel b/src/test/java/com/google/api/generator/gapic/composer/grpc/BUILD.bazel
index 825635a75e..7cce0fde60 100644
--- a/src/test/java/com/google/api/generator/gapic/composer/grpc/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/gapic/composer/grpc/BUILD.bazel
@@ -1,5 +1,5 @@
 load("@rules_java//java:defs.bzl", "java_proto_library", "java_test")
-load("//:rules_bazel/java/java_diff_test.bzl", "golden_update")
+load("//:rules_bazel/unit_golden_update.bzl", "golden_update")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -64,14 +64,14 @@ java_proto_library(
     name = test_name,
     size = "medium",
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ] + COMMON_SRCS,
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/grpc/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.grpc.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.grpc.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
 
@@ -80,15 +80,15 @@ java_proto_library(
 # `ServiceClient*` tests are not supported now since they are still in active development.
 
 [golden_update(
-    name = "{0}_update".format(test_name),
+    name = "%s_update" % test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ] + COMMON_SRCS,
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/grpc/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.grpc.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.grpc.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
diff --git a/src/test/java/com/google/api/generator/gapic/composer/resourcename/BUILD.bazel b/src/test/java/com/google/api/generator/gapic/composer/resourcename/BUILD.bazel
index c62de44173..f4fc1283f5 100644
--- a/src/test/java/com/google/api/generator/gapic/composer/resourcename/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/gapic/composer/resourcename/BUILD.bazel
@@ -1,5 +1,5 @@
 load("@rules_java//java:defs.bzl", "java_proto_library", "java_test")
-load("//:rules_bazel/java/java_diff_test.bzl", "golden_update")
+load("//:rules_bazel/unit_golden_update.bzl", "golden_update")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -13,7 +13,7 @@ TESTS = UPDATE_GOLDENS_TESTS + [
 
 filegroup(
     name = "resourcename_files",
-    srcs = ["{0}.java".format(f) for f in TESTS],
+    srcs = ["%s.java" % f for f in TESTS],
 )
 
 TEST_DEPS = [
@@ -40,13 +40,13 @@ TEST_DEPS = [
 
 [java_test(
     name = test_name,
-    srcs = ["{0}.java".format(test_name)],
+    srcs = ["%s.java" % test_name],
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/resourcename/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.resourcename.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.resourcename.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
 
@@ -62,15 +62,15 @@ java_proto_library(
 # `ServiceClient*` tests are not supported now since they are still in active development.
 
 [golden_update(
-    name = "{0}_update".format(test_name),
+    name = "%s_update" % test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ],
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/resourcename/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.resourcename.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.resourcename.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in UPDATE_GOLDENS_TESTS]
diff --git a/src/test/java/com/google/api/generator/gapic/composer/rest/BUILD.bazel b/src/test/java/com/google/api/generator/gapic/composer/rest/BUILD.bazel
index c4532b3599..ac8ecbbb5b 100644
--- a/src/test/java/com/google/api/generator/gapic/composer/rest/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/gapic/composer/rest/BUILD.bazel
@@ -1,5 +1,5 @@
 load("@rules_java//java:defs.bzl", "java_proto_library", "java_test")
-load("//:rules_bazel/java/java_diff_test.bzl", "golden_update")
+load("//:rules_bazel/unit_golden_update.bzl", "golden_update")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -62,14 +62,14 @@ java_proto_library(
     name = test_name,
     size = "medium",
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ] + COMMON_SRCS,
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/rest/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.rest.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.rest.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
 
@@ -78,15 +78,15 @@ java_proto_library(
 # `ServiceClient*` tests are not supported now since they are still in active development.
 
 [golden_update(
-    name = "{0}_update".format(test_name),
+    name = "%s_update" % test_name,
     srcs = [
-        "{0}.java".format(test_name),
+        "%s.java" % test_name,
     ] + COMMON_SRCS,
     data = [
         "//src/test/java/com/google/api/generator/gapic/composer/rest/goldens:goldens_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:gapic_config_files",
         "//src/test/java/com/google/api/generator/gapic/testdata:service_config_files",
     ],
-    test_class = "com.google.api.generator.gapic.composer.rest.{0}".format(test_name),
+    test_class = "com.google.api.generator.gapic.composer.rest.%s" % test_name,
     deps = TEST_DEPS,
 ) for test_name in TESTS]
diff --git a/src/test/java/com/google/api/generator/test/framework/BUILD.bazel b/src/test/java/com/google/api/generator/test/framework/BUILD.bazel
index c9f4b9ce74..99800b9c3f 100644
--- a/src/test/java/com/google/api/generator/test/framework/BUILD.bazel
+++ b/src/test/java/com/google/api/generator/test/framework/BUILD.bazel
@@ -27,9 +27,7 @@ java_library(
 
 java_library(
     name = "junit_runner",
-    srcs = [
-        "SingleJUnitTestRunner.java",
-    ],
+    srcs = ["SingleJUnitTestRunner.java"],
     deps = [
         ":utils",
         "@junit_junit//jar",
diff --git a/src/test/java/com/google/api/generator/test/framework/Utils.java b/src/test/java/com/google/api/generator/test/framework/Utils.java
index d38142a582..01112bf6b8 100644
--- a/src/test/java/com/google/api/generator/test/framework/Utils.java
+++ b/src/test/java/com/google/api/generator/test/framework/Utils.java
@@ -21,8 +21,7 @@
 
 public class Utils {
   /**
-   * Save the generated code from JUnit test to a file for updating goldens. These files will be
-   * saved as a zip file, then unzipped to overwrite goldens files. The relative path
+   * Save the generated code from JUnit test to a file for updating goldens. The relative path
    * `com/google/..` which is identical with the location of goldens files which will help us easily
    * replace the original goldens. For example:
    * `src/test/java/com/google/api/generator/gapic/composer/ComposerTest.java` will save the