diff --git a/.bazelrc b/.bazelrc
index ae5ba807..f51c9510 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -429,4 +429,3 @@ build:rbe_linux_cpu --python_path="/usr/bin/python3"
 # These you may need to change for your own GCP project.
 common:rbe_linux_cpu --remote_instance_name=projects/tensorflow-testing/instances/default_instance
 # END LITERT REMOTE BUILD EXECUTION OPTIONS
-
diff --git a/ci/run_bazel_test.sh b/ci/run_bazel_test.sh
index e9220a61..a89f4cf2 100755
--- a/ci/run_bazel_test.sh
+++ b/ci/run_bazel_test.sh
@@ -17,10 +17,10 @@ set -ex
 
 # Run this script under the root directory.
 
-TEST_LANG_FILTERS="${TEST_LANG_FILTERS:-cc,py}"
 EXPERIMENTAL_TARGETS_ONLY="${EXPERIMENTAL_TARGETS_ONLY:-false}"
 PUBLIC_CACHE_PUSH="${PUBLIC_CACHE_PUSH:-false}"
 PUBLIC_CACHE="${PUBLIC_CACHE:-false}"
+TEST_LANG_FILTERS="${TEST_LANG_FILTERS:-cc,py}"
 
 BUILD_FLAGS=("-c" "opt"
     "--cxxopt=--std=c++17"
@@ -47,13 +47,8 @@ BUILD_FLAGS=("-c" "opt"
     "--experimental_ui_max_stdouterr_bytes=3145728"
   )
 
-if [ "$PUBLIC_CACHE_PUSH" == "true" ]; then
-    BUILD_FLAGS+=("--config=public_cache_push")
-fi
-
-if [ "$PUBLIC_CACHE" == "true" ]; then
-    BUILD_FLAGS+=("--config=public_cache")
-fi
+# Add Bazel --config flags based on kokoro injected env ie. --config=public_cache
+BUILD_FLAGS += ( ${BAZEL_CONFIG_FLAGS} )
 
 # TODO: (b/381310257) - Investigate failing test not included in cpu_full
 # TODO: (b/381110338) - Clang errors
diff --git a/ci/run_bazel_test_with_docker.sh b/ci/run_bazel_test_with_docker.sh
index 5de21f00..d0b02b4f 100755
--- a/ci/run_bazel_test_with_docker.sh
+++ b/ci/run_bazel_test_with_docker.sh
@@ -30,7 +30,7 @@ if [ ! -d /root_dir ]; then
     -v ${SCRIPT_DIR}:/script_dir \
     -e DOCKER_PYTHON_VERSION=${DOCKER_PYTHON_VERSION} \
     -e EXPERIMENTAL_TARGETS_ONLY=${EXPERIMENTAL_TARGETS_ONLY:-false} \
-    -e PUBLIC_CACHE_PUSH=${PUBLIC_CACHE_PUSH:-false} \
+    -e BAZEL_CONFIG_FLAGS=${BAZEL_CONFIG_FLAGS} \
     --entrypoint /script_dir/run_bazel_test_with_docker.sh \
     tflite-builder
   exit 0