| 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: Keep-Alive Daemon Lifecycle |
| 17 | # Verifies that `colab new` spawns a detached keep-alive daemon, persists its |
| 18 | # PID in the session state, and that `colab stop` reaps it cleanly. |
| 19 | # |
| 20 | # This test is a fast smoke test (~10s). For a soak test that verifies the |
| 21 | # daemon's pings actually succeed against the live backend, see |
| 22 | # integration/repro_keep_alive_scope/. |
| 23 | |
| 24 | set -e |
| 25 | |
| 26 | # Setup a clean session file for testing |
| 27 | TMP_DIR=$(mktemp -d) |
| 28 | SESSION_FILE="$TMP_DIR/sessions.json" |
| 29 | trap "rm -rf $TMP_DIR" EXIT |
| 30 | |
| 31 | # To exercise the daemon we need OAuth2 or ADC with the `colaboratory` |
| 32 | # scope. Selection priority: OAuth2 (cached token present) > ADC (with the |
| 33 | # right scopes). |
| 34 | EXPECT_DAEMON=1 |
| 35 | if [ -f "$HOME/.config/colab-cli/token.json" ]; then |
| 36 | AUTH_FLAGS="--auth=oauth2" |
| 37 | elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then |
| 38 | # Check that ADC has both required scopes (userinfo.email + colaboratory). |
| 39 | ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) |
| 40 | 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) |
| 41 | if echo "$ADC_SCOPES" | grep -q "colaboratory" && echo "$ADC_SCOPES" | grep -q "userinfo.email"; then |
| 42 | AUTH_FLAGS="--auth=adc" |
| 43 | else |
| 44 | echo "Error: ADC token lacks the required scopes." |
| 45 | echo "Re-issue ADC creds with all required scopes:" |
| 46 | echo " gcloud auth application-default login \\" |
| 47 | echo " --scopes=openid,\\" |
| 48 | echo " https://www.googleapis.com/auth/cloud-platform,\\" |
| 49 | echo " https://www.googleapis.com/auth/userinfo.email,\\" |
| 50 | echo " https://www.googleapis.com/auth/colaboratory" |
| 51 | exit 1 |
| 52 | fi |
| 53 | else |
| 54 | echo "Error: No usable auth provider found." |
| 55 | echo "Options:" |
| 56 | echo " - OAuth2: run 'uv run colab --auth=oauth2 sessions' to bootstrap" |
| 57 | echo " - ADC: gcloud auth application-default login \\" |
| 58 | echo " --scopes=openid,\\" |
| 59 | echo " https://www.googleapis.com/auth/cloud-platform,\\" |
| 60 | echo " https://www.googleapis.com/auth/userinfo.email,\\" |
| 61 | echo " https://www.googleapis.com/auth/colaboratory" |
| 62 | exit 1 |
| 63 | fi |
| 64 | |
| 65 | SESSION_NAME="test-live-keep-alive" |
| 66 | |
| 67 | cleanup_session() { |
| 68 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" 2>/dev/null || true |
| 69 | } |
| 70 | trap "cleanup_session; rm -rf $TMP_DIR" EXIT |
| 71 | |
| 72 | echo "[*] Creating new session (REAL API CALL) using $AUTH_FLAGS..." |
| 73 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME" |
| 74 | |
| 75 | # Verify session exists in state |
| 76 | if [ ! -f "$SESSION_FILE" ]; then |
| 77 | echo "Error: Session file '$SESSION_FILE' not created." |
| 78 | exit 1 |
| 79 | fi |
| 80 | |
| 81 | grep "$SESSION_NAME" "$SESSION_FILE" |
| 82 | |
| 83 | # Extract PID (may be null if keep-alive was intentionally disabled). |
| 84 | PID=$(grep -A 15 "$SESSION_NAME" "$SESSION_FILE" | grep "keep_alive_pid" | awk '{print $2}' | tr -d ',') |
| 85 | |
| 86 | if [ -z "$PID" ] || [ "$PID" == "null" ]; then |
| 87 | echo "[FAILURE] No keep_alive_pid found under $AUTH_FLAGS (daemon should have spawned)." |
| 88 | cat "$SESSION_FILE" |
| 89 | exit 1 |
| 90 | fi |
| 91 | echo "[*] Keep-alive PID: $PID" |
| 92 | |
| 93 | if ps -p $PID > /dev/null; then |
| 94 | echo "[*] Keep-alive process is running." |
| 95 | else |
| 96 | echo "[FAILURE] Keep-alive process NOT running." |
| 97 | exit 1 |
| 98 | fi |
| 99 | |
| 100 | # Verify the process command line is actually colab keep-alive |
| 101 | ps -fp $PID | grep "keep-alive" |
| 102 | |
| 103 | LOG_OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" log -s "$SESSION_NAME") |
| 104 | echo "$LOG_OUTPUT" |
| 105 | if ! echo "$LOG_OUTPUT" | grep -q "KEEP: started"; then |
| 106 | echo "[FAILURE] keep_alive_started event missing from history." |
| 107 | exit 1 |
| 108 | fi |
| 109 | # The pre-flight in `colab new` calls keep_alive_assignment once |
| 110 | # synchronously before returning. If that succeeded, the structured |
| 111 | # history should NOT contain any KEEP: error events at this point. |
| 112 | if echo "$LOG_OUTPUT" | grep -q "KEEP: error"; then |
| 113 | echo "[FAILURE] keep_alive_error events present immediately after 'colab new'." |
| 114 | echo " The pre-flight keep-alive ping failed. Check the body= field." |
| 115 | exit 1 |
| 116 | fi |
| 117 | |
| 118 | echo "[*] Stopping session (REAL API CALL)..." |
| 119 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" |
| 120 | sleep 1 |
| 121 | |
| 122 | if [ "$EXPECT_DAEMON" -eq 1 ]; then |
| 123 | if ! ps -p $PID > /dev/null; then |
| 124 | echo "[*] Keep-alive process terminated successfully." |
| 125 | else |
| 126 | echo "[FAILURE] Keep-alive process still running after stop!" |
| 127 | kill $PID |
| 128 | exit 1 |
| 129 | fi |
| 130 | fi |
| 131 | |
| 132 | # Disable the cleanup trap; we already cleaned up. |
| 133 | trap "rm -rf $TMP_DIR" EXIT |
| 134 | |
| 135 | echo "[SUCCESS] Live integration test passed!" |