diff --git a/tools/upgrade.sh b/tools/upgrade.sh
index d7016aa44..5eb90ab41 100755
--- a/tools/upgrade.sh
+++ b/tools/upgrade.sh
@@ -10,9 +10,14 @@ fi
 
 # Protect against unwanted sourcing
 case "$ZSH_EVAL_CONTEXT" in
-  *:file) echo "error: this file should not be sourced" && return ;;
+  *:file) echo "error: this file should not be sourced" && return 1 ;;
 esac
 
+# Define "$ZSH" if not defined -- in theory this should be `export`ed by the calling script
+if [[ -z "$ZSH" ]]; then
+  ZSH="${0:a:h:h}"
+fi
+
 cd "$ZSH"
 
 verbose_mode="default"