tests/test_connections/test_endpoints.sh

#!/usr/bin/env bash
set -euo pipefail

# --- Safe .env load (strip CRLF if present) ---
ENV_FILE="${ENV_FILE:-.env}"
if [ -f "$ENV_FILE" ]; then
  tmp_env="$(mktemp)"
  sed 's/\r$//' "$ENV_FILE" > "$tmp_env"
  # shellcheck disable=SC1090
  . "$tmp_env"
  rm -f "$tmp_env"
fi

# Defaults (can be overridden via .env or CLI)
IMG_BASE="${AI_FAST_IMAGE_SERVER_API_GRADIO_URL:-https://ruslanmv-ai-fast-image-server.hf.space}"
IMG_TOKEN="${AI_FAST_IMAGE_SERVER_API_SECRET_TOKEN:-default_secret}"

STORY_BASE="${AI_STORY_API_GRADIO_URL:-https://ruslanmv-ai-story-server-cpu.hf.space}"
STORY_TOKEN="${AI_STORY_API_SECRET_TOKEN:-secret}"

STORY_PROMPT="${1:-A short wholesome bedtime story about a friendly dragon.}"
VOICE="${2:-Cloée}"
IMG_PROMPT="${3:-a cute cat sticker}"

# Optional Bearer (for private Space / ZeroGPU quota)
AUTH_HEADER=()
if [[ -n "${HF_TOKEN:-}" ]]; then
  AUTH_HEADER=(-H "Authorization: Bearer ${HF_TOKEN}")
fi

has_jq() { command -v jq >/dev/null 2>&1; }
pretty() { if has_jq; then jq . || cat; else cat; fi; }

# Extract JSON payload from any noisy output: keep from first { or [ to end
strip_to_json() {
  awk '
    BEGIN{ found=0 }
    {
      if (!found) {
        pos = index($0, "{"); if (pos==0) pos = index($0, "[");
        if (pos>0) { print substr($0, pos); found=1; }
      } else { print $0; }
    }'
}

# POST helper (returns BODY + trailing HTTP_STATUS line)
post_json() {
  local url="$1" json="$2"
  echo "POST $url" >&2
  curl -sS -H "Content-Type: application/json" "${AUTH_HEADER[@]}" \
       -X POST "$url" -d "$json" -w "\nHTTP_STATUS:%{http_code}"
}

# GET helper (returns BODY + trailing HTTP_STATUS line)
get_url() {
  local url="$1"
  echo "GET  $url" >&2
  curl -sS "${AUTH_HEADER[@]}" "$url" -w "\nHTTP_STATUS:%{http_code}"
}

# --- Parse a Gradio SSE GET stream and echo the `complete` JSON object ---
# Returns 0 on success and prints the JSON object to stdout.
# Returns 1 on SSE `error` (prints details to stderr if present).
# Returns 2 on malformed/unknown stream.
parse_sse_complete() {
  local stream_file="$1"
  local event="" data_buf="" complete_json="" error_json=""

  # Normalize CRLF in stream (if any)
  sed -i 's/\r$//' "$stream_file"

  while IFS= read -r line || [[ -n "$line" ]]; do
    if [[ "$line" == event:* ]]; then
      event="${line#event: }"
      continue
    fi
    if [[ "$line" == data:* ]]; then
      local d="${line#data: }"
      if [[ -z "$data_buf" ]]; then data_buf="$d"; else data_buf+=$'\n'"$d"; fi
      continue
    fi
    # blank line = event separator
    if [[ -z "$line" ]]; then
      if [[ "$event" == "complete" ]]; then
        complete_json="$data_buf"; break
      elif [[ "$event" == "error" ]]; then
        error_json="$data_buf"; break
      fi
      event=""; data_buf=""
    fi
  done < "$stream_file"

  if [[ -n "$complete_json" ]]; then
    printf '%s' "$complete_json"
    return 0
  fi
  if [[ -n "$error_json" && "$error_json" != "null" ]]; then
    echo "SSE error event:" >&2
    echo "$error_json" >&2
    return 1
  fi

  echo "SSE stream had no complete/error event. Raw stream follows:" >&2
  cat "$stream_file" >&2
  return 2
}

# New Gradio v4/v5 call API: POST -> event_id -> GET SSE -> complete JSON
try_call_api() {
  local base="$1" func="$2" json="$3"
  local call_url="${base%/}/gradio_api/call/$func"
  local out status body body_json event_id get_url_full stream_file comp_json

  out="$(post_json "$call_url" "$json" || true)"
  status="${out##*HTTP_STATUS:}"
  body="${out%HTTP_STATUS:*}"
  body_json="$(printf '%s' "$body" | strip_to_json)"

  echo "---- ${func}: POST response ----"
  printf '%s\n' "$body_json" | pretty
  echo "--------------------------------"

  if [[ "$status" != "200" ]]; then
    echo "-> ${status} from $call_url" >&2
    printf '%s\n' "$body" | head -c 500 >&2; echo >&2
    return 1
  fi

  # Parse event_id
  if has_jq; then
    event_id="$(printf '%s' "$body_json" | jq -r '.event_id // .eventId // .eventID // empty' 2>/dev/null || echo "")"
  else
    event_id="$(printf '%s' "$body_json" | sed -n 's/.*"event_id"[[:space:]]*:[[:space:]]*"\([^"]\+\)".*/\1/p')"
  fi
  if [[ -z "${event_id:-}" ]]; then
    echo "-> Missing event_id in response." >&2
    return 1
  fi

  # GET SSE
  get_url_full="${base%/}/gradio_api/call/$func/$event_id"
  stream_file="$(mktemp)"
  echo "GET  $get_url_full" >&2
  curl -sS -N -H "Accept: text/event-stream" "${AUTH_HEADER[@]}" \
       "$get_url_full" > "$stream_file"

  # Parse SSE to the complete event JSON object
  if ! comp_json="$(parse_sse_complete "$stream_file")"; then
    rm -f "$stream_file"
    return 1
  fi
  rm -f "$stream_file"

  # Show full `complete` JSON (wrapper includes {"data":[ ... ], ...})
  echo "---- ${func}: GET (complete event) ----"
  printf '%s\n' "$comp_json" | pretty
  echo "--------------------------------------"

  # If this is the story function, also print ONLY the story component (.data[0])
  if [[ "$base" == "$STORY_BASE" && "$func" == "predict" ]]; then
    if has_jq && jq -e '.data | type == "array" and (.data | length) > 0' >/dev/null 2>&1 <<<"$comp_json"; then
      echo "---- story data (.data[0]) ----"
      jq -r '.data[0]' <<<"$comp_json"
      echo "-------------------------------"
    fi
  fi

  echo "-> 200 OK (call API)"
  return 0
}

# Old Gradio endpoints (sync JSON): /api/predict then /run/predict
try_legacy_predict() {
  local base="$1" json="$2"
  local primary="${base%/}/api/predict"
  local fallback="${base%/}/run/predict"
  local out status body body_json

  out="$(post_json "$primary" "$json" || true)"
  status="${out##*HTTP_STATUS:}"
  body="${out%HTTP_STATUS:*}"
  body_json="$(printf '%s' "$body" | strip_to_json)"

  echo "---- legacy POST ($primary) response ----"
  printf '%s\n' "$body_json" | pretty
  echo "-----------------------------------------"

  if [[ "$status" == "200" ]]; then
    echo "-> 200 OK ($primary)"
    return 0
  else
    echo "-> ${status} from $primary" >&2
    printf '%s\n' "$body" | head -c 500 >&2; echo >&2
  fi

  out="$(post_json "$fallback" "$json" || true)"
  status="${out##*HTTP_STATUS:}"
  body="${out%HTTP_STATUS:*}"
  body_json="$(printf '%s' "$body" | strip_to_json)"

  echo "---- legacy POST ($fallback) response ----"
  printf '%s\n' "$body_json" | pretty
  echo "------------------------------------------"

  if [[ "$status" == "200" ]]; then
    echo "-> 200 OK ($fallback)"
    return 0
  else
    echo "-> ${status} from $fallback" >&2
    printf '%s\n' "$body" | head -c 500 >&2; echo >&2
    return 1
  fi
}

# ---------- Build JSON payloads ----------
# Image server /generate expects:
# [prompt, negative_prompt, seed, width, height, guidance, steps, secret_token]
IMG_JSON=$(cat <<EOF
{"data": [
  "${IMG_PROMPT}",
  "",
  3,
  256,
  256,
  0,
  1,
  "${IMG_TOKEN}"
]}
EOF
)

# Story server /predict expects:
# [secret_token, story_prompt, voice]
STORY_JSON=$(cat <<EOF
{"data": [
  "${STORY_TOKEN}",
  "${STORY_PROMPT}",
  "${VOICE}"
]}
EOF
)

echo "=== Image server check ==="
echo "Base: $IMG_BASE"
if ! try_call_api "$IMG_BASE" "generate" "$IMG_JSON"; then
  echo
  echo "Falling back to legacy endpoints for image server…"
  if ! try_legacy_predict "$IMG_BASE" "{\"fn_index\":0, \"data\": [\"${IMG_PROMPT}\", \"\", 3, 256, 256, 0, 1, \"${IMG_TOKEN}\"]}"; then
    echo "Image endpoint check FAILED" >&2
    exit 2
  fi
fi

echo
echo "=== Story server check ==="
echo "Base: $STORY_BASE"
if ! try_call_api "$STORY_BASE" "predict" "$STORY_JSON"; then
  echo
  echo "Falling back to legacy endpoints for story server…"
  if ! try_legacy_predict "$STORY_BASE" "{\"fn_index\":0, \"data\": [\"${STORY_TOKEN}\", \"${STORY_PROMPT}\", \"${VOICE}\"]}"; then
    echo "Story endpoint check FAILED" >&2
    exit 3
  fi
fi

echo
echo "All endpoint checks completed successfully."