From d06eedae86214a6f5cddf47750ebabeff1bf0d3d Mon Sep 17 00:00:00 2001
From: "Jorge E. Moreira" <jemoreira@google.com>
Date: Mon, 25 Mar 2024 16:20:17 -0700
Subject: [PATCH] Allow passing credential sources to e2e tests

---
 e2etests/boot_tests.bzl          | 14 ++++++++++----
 e2etests/cvd_load_boot_test.sh   | 31 +++++++++++++++++++++++++------
 e2etests/launch_cvd_boot_test.sh | 11 +++++++++--
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/e2etests/boot_tests.bzl b/e2etests/boot_tests.bzl
index 9ac6b3368d..4cd14179e3 100644
--- a/e2etests/boot_tests.bzl
+++ b/e2etests/boot_tests.bzl
@@ -1,9 +1,12 @@
-def launch_cvd_boot_test(name, branch, target):
+def launch_cvd_boot_test(name, branch, target, credential_source = ""):
+    args = ["-b", branch, "-t", target]
+    if credential_source:
+        args += ["-c", credential_source]
     native.sh_test(
         name = name,
         size = "medium",
         srcs = ["launch_cvd_boot_test.sh"],
-        args = ["-b", branch, "-t", target],
+        args = args,
         tags = [
             "exclusive",
             "external",
@@ -11,12 +14,15 @@ def launch_cvd_boot_test(name, branch, target):
         ],
     )
 
-def cvd_load_boot_test(name, env_file, size = "medium"):
+def cvd_load_boot_test(name, env_file, size = "medium", credential_source = ""):
+    args = ["-e", env_file]
+    if credential_source:
+        args += ["-c", credential_source]
     native.sh_test(
         name = name,
         size = size,
         srcs = ["cvd_load_boot_test.sh"],
-        args = [env_file],
+        args = args,
         data = [env_file],
         tags = [
             "exclusive",
diff --git a/e2etests/cvd_load_boot_test.sh b/e2etests/cvd_load_boot_test.sh
index 90176cd86a..2eeaab7920 100755
--- a/e2etests/cvd_load_boot_test.sh
+++ b/e2etests/cvd_load_boot_test.sh
@@ -2,14 +2,29 @@
 
 set -e -x
 
-if [[ $# != 1 ]]; then
-  echo "Usage: $0 ENVIRONMENT_SPECIFICATION_FILE"
+usage="Usage: $0 [-c CREDENTIAL_SOURCE] ENVIRONMENT_SPECIFICATION_FILE"
+
+ENV_FILE=""
+CREDENTIAL_SOURCE=""
+while getopts "c:e:" opt; do
+  case "${opt}" in
+    c)
+      CREDENTIAL_SOURCE="${OPTARG}"
+      ;;
+    e)
+      ENV_FILE="${OPTARG}"
+      ;;
+    *)
+      echo "${usage}"
+      exit 1
+  esac
+done
+
+if [[ -z "${ENV_FILE}" ]]; then
+  echo "${usage}"
   exit 1
 fi
 
-ls
-
-ENV_FILE="$1"
 CMD_OUT="cvd_load_stdout.txt"
 CMD_ERR="cvd_load_stderr.txt"
 
@@ -45,4 +60,8 @@ trap collect_logs_and_cleanup EXIT
 # Make sure there is no cvd server around
 cvd reset -y
 
-cvd load "${ENV_FILE}" >"${CMD_OUT}" 2>"${CMD_ERR}"
+credential_arg=""
+if [[ -n "$CREDENTIAL_SOURCE" ]]; then
+    credential_arg="--override=fetch.credential_source:${CREDENTIAL_SOURCE}"
+fi
+cvd load ${credential_arg} "${ENV_FILE}" >"${CMD_OUT}" 2>"${CMD_ERR}"
diff --git a/e2etests/launch_cvd_boot_test.sh b/e2etests/launch_cvd_boot_test.sh
index 0b12ec4971..102d7452e9 100755
--- a/e2etests/launch_cvd_boot_test.sh
+++ b/e2etests/launch_cvd_boot_test.sh
@@ -4,12 +4,16 @@ set -e -x
 
 BRANCH=""
 TARGET=""
+CREDENTIAL_SOURCE=""
 
-while getopts "b:t:" opt; do
+while getopts "c:b:t:" opt; do
   case "${opt}" in
     b)
       BRANCH="${OPTARG}"
       ;;
+    c)
+      CREDENTIAL_SOURCE="${OPTARG}"
+      ;;
     t)
       TARGET="${OPTARG}"
       ;;
@@ -51,7 +55,10 @@ trap collect_logs_and_cleanup EXIT
 # client can still connect to the server outside the sandbox and cause issues.
 cvd reset -y
 
-cvd fetch --default_build="${BRANCH}/${TARGET}" --target_directory="${workdir}"
+cvd fetch \
+  --default_build="${BRANCH}/${TARGET}" \
+  --target_directory="${workdir}" \
+  --credential_source="${CREDENTIAL_SOURCE}"
 
 (
   cd "${workdir}"