diff --git a/src/test/shell/integration/BUILD b/src/test/shell/integration/BUILD
index 1d9dd96..62faad4 100644
--- a/src/test/shell/integration/BUILD
+++ b/src/test/shell/integration/BUILD
@@ -408,11 +408,13 @@
     name = "nonincremental_builds_test",
     size = "medium",
     srcs = ["nonincremental_builds_test.sh"],
+    args = ["$(JAVABASE)"],
     data = [
         ":discard_graph_edges_lib.sh",
         ":test-deps",
         "@bazel_tools//tools/bash/runfiles",
     ],
+    toolchains = ["@bazel_tools//tools/jdk:current_java_runtime"],
 )
 
 sh_test(
diff --git a/src/test/shell/integration/nonincremental_builds_test.sh b/src/test/shell/integration/nonincremental_builds_test.sh
index cba0744..e55d521 100755
--- a/src/test/shell/integration/nonincremental_builds_test.sh
+++ b/src/test/shell/integration/nonincremental_builds_test.sh
@@ -16,17 +16,16 @@
 #
 # nonincremental_builds_test.sh: tests for the --keep_state_after_build flag.
 
-# --- begin runfiles.bash initialization ---
-# Copy-pasted from Bazel's Bash runfiles library (tools/bash/runfiles/runfiles.bash).
 set -euo pipefail
+# --- begin runfiles.bash initialization ---
 if [[ ! -d "${RUNFILES_DIR:-/dev/null}" && ! -f "${RUNFILES_MANIFEST_FILE:-/dev/null}" ]]; then
-  if [[ -f "$0.runfiles_manifest" ]]; then
-    export RUNFILES_MANIFEST_FILE="$0.runfiles_manifest"
-  elif [[ -f "$0.runfiles/MANIFEST" ]]; then
-    export RUNFILES_MANIFEST_FILE="$0.runfiles/MANIFEST"
-  elif [[ -f "$0.runfiles/bazel_tools/tools/bash/runfiles/runfiles.bash" ]]; then
-    export RUNFILES_DIR="$0.runfiles"
-  fi
+    if [[ -f "$0.runfiles_manifest" ]]; then
+      export RUNFILES_MANIFEST_FILE="$0.runfiles_manifest"
+    elif [[ -f "$0.runfiles/MANIFEST" ]]; then
+      export RUNFILES_MANIFEST_FILE="$0.runfiles/MANIFEST"
+    elif [[ -f "$0.runfiles/bazel_tools/tools/bash/runfiles/runfiles.bash" ]]; then
+      export RUNFILES_DIR="$0.runfiles"
+    fi
 fi
 if [[ -f "${RUNFILES_DIR:-/dev/null}/bazel_tools/tools/bash/runfiles/runfiles.bash" ]]; then
   source "${RUNFILES_DIR}/bazel_tools/tools/bash/runfiles/runfiles.bash"
@@ -64,6 +63,18 @@
   export MSYS2_ARG_CONV_EXCL="*"
 fi
 
+if "$is_windows"; then
+  EXE_EXT=".exe"
+else
+  EXE_EXT=""
+fi
+
+javabase="$1"
+if [[ $javabase = external/* ]]; then
+  javabase=${javabase#external/}
+fi
+jmaptool="$(rlocation "${javabase}/bin/jmap${EXE_EXT}")"
+
 if ! type try_with_timeout >&/dev/null; then
   # Bazel's testenv.sh defines try_with_timeout but the Google-internal version
   # uses a different testenv.sh.
@@ -129,7 +140,7 @@
     bazel build $pkg:top &> "$TEST_log"  \
         || fail "Couldn't build $pkg"
     local server_pid="$(bazel info server_pid 2>> "$TEST_log")"
-    "$bazel_javabase"/bin/jmap -histo:live "$server_pid" > histo.txt
+    "$jmaptool" -histo:live "$server_pid" > histo.txt
 
     cat histo.txt >> "$TEST_log"
     assert_contains "GenRuleAction" histo.txt
@@ -143,7 +154,7 @@
     bazel build --nokeep_state_after_build $pkg:top &> "$TEST_log"  \
         || fail "Couldn't build $pkg"
     local server_pid="$(bazel info server_pid 2>> "$TEST_log")"
-    "$bazel_javabase"/bin/jmap -histo:live "$server_pid" > histo.txt
+    "$jmaptool" -histo:live "$server_pid" > histo.txt
 
     cat histo.txt >> "$TEST_log"
     assert_not_contains "GenRuleAction$" histo.txt