| 1 | #!/bin/bash |
| 2 | # Copyright 2026 Google LLC |
| 3 | # |
| 4 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | # you may not use this file except in compliance with the License. |
| 6 | # You may obtain a copy of the License at |
| 7 | # |
| 8 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | # |
| 10 | # Unless required by applicable law or agreed to in writing, software |
| 11 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | # See the License for the specific language governing permissions and |
| 14 | # limitations under the License. |
| 15 | |
| 16 | # Integration Test: `colab run <script.py> [args...]` |
| 17 | # |
| 18 | # Verifies the shebang-friendly one-shot execution flow: |
| 19 | # 1. `colab run` allocates a CPU VM, runs the script, releases the VM. |
| 20 | # 2. `sys.argv` and `__name__ == "__main__"` are honored. |
| 21 | # 3. After the run finishes, no orphan VMs remain. |
| 22 | # 4. `colab run --keep` leaves the session alive; `colab stop` clears it. |
| 23 | |
| 24 | # Don't `set -e` so we can capture failures and clean up explicitly. |
| 25 | |
| 26 | # ---------- Auth detection (mirrors integration/repro_keep_alive/test.sh) ---- |
| 27 | if [ -f "$HOME/.config/colab-cli/token.json" ]; then |
| 28 | AUTH_FLAGS="--auth=oauth2" |
| 29 | elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then |
| 30 | ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) |
| 31 | ADC_SCOPES=$(curl -s "https://www.googleapis.com/oauth2/v3/tokeninfo?access_token=$ADC_TOKEN" | python3 -c "import json,sys; print(json.load(sys.stdin).get('scope',''))" 2>/dev/null) |
| 32 | if echo "$ADC_SCOPES" | grep -q "colaboratory" && echo "$ADC_SCOPES" | grep -q "userinfo.email"; then |
| 33 | AUTH_FLAGS="--auth=adc" |
| 34 | else |
| 35 | echo "Error: ADC token lacks the required scopes (colaboratory + userinfo.email)." |
| 36 | echo "Re-issue ADC creds with all required scopes:" |
| 37 | echo " gcloud auth application-default login \\" |
| 38 | echo " --scopes=openid,\\" |
| 39 | echo " https://www.googleapis.com/auth/cloud-platform,\\" |
| 40 | echo " https://www.googleapis.com/auth/userinfo.email,\\" |
| 41 | echo " https://www.googleapis.com/auth/colaboratory" |
| 42 | exit 1 |
| 43 | fi |
| 44 | else |
| 45 | echo "Error: No usable auth provider found." |
| 46 | exit 1 |
| 47 | fi |
| 48 | echo "[*] Using $AUTH_FLAGS" |
| 49 | |
| 50 | # ---------- Isolated session state ------------------------------------------- |
| 51 | TMP_DIR=$(mktemp -d) |
| 52 | SESSION_FILE="$TMP_DIR/sessions.json" |
| 53 | SCRIPT_PATH="$TMP_DIR/script.py" |
| 54 | KEEP_SESSION_NAME="repro-run-keep-$(date +%s)" |
| 55 | |
| 56 | cleanup() { |
| 57 | echo "[*] Cleaning up..." |
| 58 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$KEEP_SESSION_NAME" 2>/dev/null || true |
| 59 | rm -rf "$TMP_DIR" |
| 60 | } |
| 61 | trap cleanup EXIT |
| 62 | |
| 63 | cat > "$SCRIPT_PATH" <<'PYEOF' |
| 64 | import sys |
| 65 | print(f"argv={sys.argv}") |
| 66 | print(f"is_main={__name__ == '__main__'}") |
| 67 | PYEOF |
| 68 | SCRIPT_BASENAME=$(basename "$SCRIPT_PATH") |
| 69 | |
| 70 | # ---------- Phase 1: basic run + auto-cleanup -------------------------------- |
| 71 | echo "[*] Phase 1: colab run <script.py> hello world" |
| 72 | OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" run "$SCRIPT_PATH" hello world 2>&1) |
| 73 | RC=$? |
| 74 | echo "$OUTPUT" |
| 75 | |
| 76 | if [ $RC -ne 0 ]; then |
| 77 | echo "[FAILURE] colab run exited $RC" |
| 78 | exit 1 |
| 79 | fi |
| 80 | if ! echo "$OUTPUT" | grep -q "argv=\['$SCRIPT_BASENAME', 'hello', 'world'\]"; then |
| 81 | echo "[FAILURE] argv was not propagated as expected." |
| 82 | echo " Wanted substring: argv=['$SCRIPT_BASENAME', 'hello', 'world']" |
| 83 | exit 1 |
| 84 | fi |
| 85 | if ! echo "$OUTPUT" | grep -q "is_main=True"; then |
| 86 | echo "[FAILURE] __name__ was not set to '__main__'." |
| 87 | exit 1 |
| 88 | fi |
| 89 | |
| 90 | # Verify cleanup actually happened — no orphan assignments remain. |
| 91 | SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1) |
| 92 | echo "$SESSIONS_OUT" |
| 93 | if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then |
| 94 | echo "[FAILURE] After auto-cleanup, server still reports active sessions." |
| 95 | echo " (Possible orphan VM — investigate.)" |
| 96 | exit 1 |
| 97 | fi |
| 98 | echo "[SUCCESS] Phase 1 passed: argv passthrough, __main__, and auto-cleanup." |
| 99 | |
| 100 | # ---------- Phase 2: --keep leaves the session alive ------------------------- |
| 101 | echo "" |
| 102 | echo "[*] Phase 2: colab run --keep -s $KEEP_SESSION_NAME <script.py>" |
| 103 | OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" run --keep -s "$KEEP_SESSION_NAME" "$SCRIPT_PATH" keep_arg 2>&1) |
| 104 | RC=$? |
| 105 | echo "$OUTPUT" |
| 106 | |
| 107 | if [ $RC -ne 0 ]; then |
| 108 | echo "[FAILURE] colab run --keep exited $RC" |
| 109 | exit 1 |
| 110 | fi |
| 111 | if ! echo "$OUTPUT" | grep -q "argv=\['$SCRIPT_BASENAME', 'keep_arg'\]"; then |
| 112 | echo "[FAILURE] --keep run did not produce the expected argv output." |
| 113 | exit 1 |
| 114 | fi |
| 115 | |
| 116 | SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1) |
| 117 | echo "$SESSIONS_OUT" |
| 118 | if ! echo "$SESSIONS_OUT" | grep -q "\[$KEEP_SESSION_NAME\]"; then |
| 119 | echo "[FAILURE] --keep session $KEEP_SESSION_NAME not found in colab sessions." |
| 120 | exit 1 |
| 121 | fi |
| 122 | |
| 123 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$KEEP_SESSION_NAME" |
| 124 | SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1) |
| 125 | if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then |
| 126 | echo "[FAILURE] After manual stop of $KEEP_SESSION_NAME, sessions remain." |
| 127 | exit 1 |
| 128 | fi |
| 129 | |
| 130 | echo "[SUCCESS] Phase 2 passed: --keep persists the session, manual stop clears it." |
| 131 | echo "[SUCCESS] All phases passed." |
| 132 | exit 0 |