refactor, reset on shutdown

mafredri · mafredri · commit a277b4f5e7e9 · 2023-09-29T17:24:54.000Z
diff --git a/scaletest/templates/scaletest-runner/scripts/lib.sh b/scaletest/templates/scaletest-runner/scripts/lib.sh
@@ -210,3 +210,33 @@ wait_baseline() {
 	maybedryrun "$DRY_RUN" sleep $((s * 60))
 	PHASE_TYPE="phase-wait" end_phase
 }
+
+get_appearance() {
+	session_token=$CODER_SESSION_TOKEN
+	if [[ -f "${CODER_CONFIG_DIR}/session" ]]; then
+		session_token="$(<"${CODER_CONFIG_DIR}/session")"
+	fi
+	curl -sSL \
+		-H "Coder-Session-Token: ${session_token}" \
+		"${CODER_URL}/api/v2/appearance"
+}
+set_appearance() {
+	local json=$1 color=$2 message=$3
+
+	session_token=$CODER_SESSION_TOKEN
+	if [[ -f "${CODER_CONFIG_DIR}/session" ]]; then
+		session_token="$(<"${CODER_CONFIG_DIR}/session")"
+	fi
+	newjson="$(
+		jq \
+			--arg color "${color}" \
+			--arg message "${message}" \
+			'. | .service_banner.message |= $message | .service_banner.background_color |= $color' <<<"${json}"
+	)"
+	maybedryrun "${DRY_RUN}" curl -sSL \
+		-X PUT \
+		-H 'Content-Type: application/json' \
+		-H "Coder-Session-Token: ${session_token}" \
+		--data "${newjson}" \
+		"${CODER_URL}/api/v2/appearance"
+}
diff --git a/scaletest/templates/scaletest-runner/shutdown.sh b/scaletest/templates/scaletest-runner/shutdown.sh
@@ -16,3 +16,10 @@ annotate_grafana "workspace" "Agent stopping..."
 "${SCRIPTS_DIR}/cleanup.sh" shutdown
 
 annotate_grafana_end "workspace" "Agent running"
+
+appearance_json="$(get_appearance)"
+service_banner_message=$(jq -r '.service_banner.message' <<<"${appearance_json}")
+service_banner_message="${service_banner_message/% | */}"
+service_banner_color="#4CD473" # Green.
+
+set_appearance "${service_banner_color}" "${service_banner_message}"
diff --git a/scaletest/templates/scaletest-runner/startup.sh b/scaletest/templates/scaletest-runner/startup.sh
@@ -21,26 +21,11 @@ fi
 # shellcheck disable=SC2153 source=scaletest/templates/scaletest-runner/scripts/lib.sh
 . "${SCRIPTS_DIR}/lib.sh"
 
-appearance_json="$(
-	curl -sSL \
-		-H "Coder-Session-Token: ${CODER_USER_TOKEN}" \
-		"${CODER_URL}/api/v2/appearance"
-)"
+appearance_json="$(get_appearance)"
 service_banner_message=$(jq -r '.service_banner.message' <<<"${appearance_json}")
 service_banner_message="${service_banner_message/% | */}"
 service_banner_color="#D65D0F" # Orange.
 
-set_appearance() {
-	local color=$1 message=$2
-	jq --arg color "${color}" --arg message "${message}" '. | .service_banner.message |= $message | .service_banner.background_color |= $color' <<<"${appearance_json}" \
-		| curl -sSL \
-			-X 'PUT' \
-			-H 'Content-Type: application/json' \
-			-H "Coder-Session-Token: ${CODER_USER_TOKEN}" \
-			--data @- \
-			"${CODER_URL}/api/v2/appearance"
-}
-
 annotate_grafana "workspace" "Agent running" # Ended in shutdown.sh.
 
 {
