This is an automated email from the ASF dual-hosted git repository. wenjin272 pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 432e90d875b7174d84d06ba8eebbd2bf0144ae51 Author: WenjinXie <[email protected]> AuthorDate: Thu Aug 6 22:46:40 2026 +0800 [e2e] Stabilize standalone example tests Generated-by: Codex CLI 0.144.5 --- .github/workflows/nightly-e2e.yml | 15 +- .../test-scripts/test_submit_examples_to_flink.sh | 562 +++++++++++++++++---- .../agents/examples/agents/ParallelChatAgent.java | 3 +- .../quickstart/agents/parallel_chat_agent.py | 9 +- .../flink_agents/examples/rag/agents/__init__.py | 17 + .../{rag_agent_example.py => agents/rag_agent.py} | 64 +-- .../examples/rag/knowledge_base_setup.py | 13 +- .../flink_agents/examples/rag/rag_agent_example.py | 142 +----- .../examples/rag/tests/test_rag_agent.py | 46 ++ tools/test/unit/verify_example_job.bats | 466 +++++++++++++++++ 10 files changed, 1037 insertions(+), 300 deletions(-) diff --git a/.github/workflows/nightly-e2e.yml b/.github/workflows/nightly-e2e.yml index 1adf197e..9863dfd3 100644 --- a/.github/workflows/nightly-e2e.yml +++ b/.github/workflows/nightly-e2e.yml @@ -23,9 +23,9 @@ on: workflow_dispatch: inputs: flink-version: - description: 'Flink version to test (e.g., 2.2.0, 2.1.1)' + description: 'Flink version to test (e.g., 2.3.0, 2.2.1)' required: false - default: '2.2.0' + default: '2.3.0' type: string pull_request: branches: [main] @@ -47,14 +47,14 @@ jobs: strategy: fail-fast: false matrix: - flink-version: ['2.2.0'] + flink-version: ['2.3.0'] steps: - uses: actions/checkout@v4 - name: Install java uses: actions/setup-java@v4 with: java-version: '11' - distribution: 'adopt' + distribution: 'temurin' - name: Install python uses: actions/setup-python@v4 with: @@ -66,9 +66,10 @@ jobs: - name: Run submit-examples E2E test env: FLINK_VERSION: ${{ github.event.inputs.flink-version || matrix.flink-version }} - # Use a lightweight chat model on GitHub-hosted runners (CPU only, - # ~7GB RAM). The script aliases this to the names examples hardcode. - OLLAMA_CHAT_MODEL: 'qwen3:0.6b' + # Use a compact model that reliably follows the examples' structured + # output prompts. The script aliases it to the names examples hardcode. + OLLAMA_CHAT_MODEL: 'qwen3.5:4b' + OLLAMA_CONTEXT_LENGTH: '2048' OLLAMA_EMBED_MODEL: 'nomic-embed-text' run: bash e2e-test/test-scripts/test_submit_examples_to_flink.sh - name: Upload Flink logs on failure diff --git a/e2e-test/test-scripts/test_submit_examples_to_flink.sh b/e2e-test/test-scripts/test_submit_examples_to_flink.sh index 68cadb04..7b479c42 100755 --- a/e2e-test/test-scripts/test_submit_examples_to_flink.sh +++ b/e2e-test/test-scripts/test_submit_examples_to_flink.sh @@ -17,12 +17,14 @@ # limitations under the License. # # -# Submits all Java/Python examples to a local Flink standalone cluster, runs -# them against a local Ollama server and waits for each job to FINISH. -# Examples are auto-discovered from the examples directories. +# Submits all Java/Python examples to a local Flink standalone cluster and runs +# them against a local Ollama server. Detached jobs pass after they either +# finish or remain RUNNING for a bounded stability period. Examples are +# auto-discovered from the examples directories. # -# Env: FLINK_VERSION (default 2.2.0), FLINK_HOME (reuse existing install), -# VERBOSE=1 (set -x). +# Env: FLINK_VERSION (default 2.3.0), FLINK_HOME (reuse existing install), +# JOB_STARTUP_TIMEOUT (default 60), JOB_STABLE_RUNNING_SECONDS (default 20), +# JOB_STATUS_POLL_INTERVAL (default 2), VERBOSE=1 (set -x). set -euo pipefail @@ -50,10 +52,14 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" ROOT_DIR="$(cd "$SCRIPT_DIR/../.."; pwd)" log_info "Project root: $ROOT_DIR" -FLINK_VERSION="${FLINK_VERSION:-2.2.0}" +FLINK_VERSION="${FLINK_VERSION:-2.3.0}" FLINK_MAJOR_MINOR="${FLINK_VERSION%.*}" SUBMIT_TIMEOUT="${SUBMIT_TIMEOUT:-180}" JOB_FINISH_TIMEOUT="${JOB_FINISH_TIMEOUT:-300}" +JOB_STARTUP_TIMEOUT="${JOB_STARTUP_TIMEOUT:-60}" +JOB_STABLE_RUNNING_SECONDS="${JOB_STABLE_RUNNING_SECONDS:-20}" +JOB_STATUS_POLL_INTERVAL="${JOB_STATUS_POLL_INTERVAL:-2}" +SLOT_RELEASE_TIMEOUT="${SLOT_RELEASE_TIMEOUT:-30}" # Models to pull for Ollama. Override these in CI to use lighter models. # The script aliases the pulled chat model to the names hardcoded in @@ -63,12 +69,13 @@ OLLAMA_EMBED_MODEL="${OLLAMA_EMBED_MODEL:-nomic-embed-text}" # Model names referenced (hardcoded) by example code. The pulled chat model # is aliased to each of these via `ollama cp` when it differs. -OLLAMA_CHAT_MODEL_ALIASES=("qwen3:8b" "qwen3.5:9b") +OLLAMA_CHAT_MODEL_ALIASES=("qwen3:1.7b" "qwen3:8b" "qwen3.5:9b") # Bash 3 (default on macOS) lacks associative arrays. RESULT_NAMES=() RESULT_STATES=() SUBMITTED_JOB_IDS=() +SUBMISSION_PIDS=() OLLAMA_PID="" cleanup() { @@ -82,6 +89,18 @@ cleanup() { "$FLINK_HOME/bin/flink" cancel "$jid" >/dev/null 2>&1 || true done + # Attached Java submissions are kept alive while their jobs are + # checked so example shutdown hooks do not remove input files before + # TaskManagers open them. Stop any client still left after cancelling + # its job. + for pid in "${SUBMISSION_PIDS[@]:-}"; do + [[ -n "$pid" ]] || continue + if kill -0 "$pid" 2>/dev/null; then + kill "$pid" 2>/dev/null || true + fi + wait "$pid" 2>/dev/null || true + done + if [[ -x "$FLINK_HOME/bin/stop-cluster.sh" ]]; then log_info "Stopping Flink cluster" "$FLINK_HOME/bin/stop-cluster.sh" >/dev/null 2>&1 || true @@ -100,7 +119,6 @@ cleanup() { print_summary exit "$exit_code" } -trap cleanup EXIT print_summary() { log_section "Test summary" @@ -126,7 +144,7 @@ print_summary() { printf "\nTotal: %d Passed: %d Failed: %d\n" "$total" "$passed" "$failed" if (( failed > 0 )); then - log_error "$failed example(s) failed to submit" + log_error "$failed example(s) failed validation" exit 1 fi } @@ -136,42 +154,95 @@ record_result() { RESULT_STATES+=("$2") } +install_flink_distribution() { + local install_dir="$1" + + # Reuse install.sh's download, archive validation, and extraction logic, + # but stop before it installs a released Flink Agents JAR. This E2E test + # must run the artifacts built from the current checkout, which may support + # a newer Flink version than the latest published Flink Agents release. + FLINK_AGENTS_INSTALL_SH_NO_RUN=1 \ + FLINK_VERSION="$FLINK_VERSION" \ + INSTALL_FLINK=Yes \ + INSTALL_DIR="$install_dir" \ + NO_PROMPT=1 \ + bash -c ' + source "$1" + plan_flink + install_flink_if_needed + ' _ "$ROOT_DIR/tools/install.sh" +} + +prepare_python_venv() { + if [[ -f "$VENV_DIR/pyvenv.cfg" && -x "$VENV_DIR/bin/python" ]]; then + log_info "Reusing Python venv: $VENV_DIR" + return + fi + if [[ -e "$VENV_DIR" ]] \ + && [[ ! -d "$VENV_DIR" \ + || -n "$(find "$VENV_DIR" -mindepth 1 -print -quit 2>/dev/null)" ]]; then + log_error "VENV_DIR is not an empty directory or valid venv: $VENV_DIR" + return 1 + fi + + local python_bin="${PYTHON_BIN:-python3}" + log_info "Creating Python venv: $VENV_DIR" + "$python_bin" -m venv "$VENV_DIR" +} + install_flink() { log_section "Step 1: install Flink standalone (version $FLINK_VERSION)" + # Anchor VENV_DIR to the repo so both fresh and reused Flink installations + # use the same Python environment. + export VENV_DIR="${VENV_DIR:-$ROOT_DIR/.flink-agents-env}" + if [[ -n "${FLINK_HOME:-}" && -x "$FLINK_HOME/bin/flink" ]]; then log_info "Reusing existing FLINK_HOME: $FLINK_HOME" export FLINK_HOME - return 0 + else + local install_dir="${INSTALL_DIR:-$HOME/.local/flink}" + log_info "Installing the Flink distribution with tools/install.sh helpers" + install_flink_distribution "$install_dir" + export FLINK_HOME="${install_dir}/flink-${FLINK_VERSION}" + + if [[ ! -x "$FLINK_HOME/bin/flink" ]]; then + log_error "Flink installation not found at expected path: $FLINK_HOME" + return 1 + fi + log_ok "Flink installed at: $FLINK_HOME" fi - # Anchor VENV_DIR to the repo so we can find it after install.sh exits. - export VENV_DIR="${VENV_DIR:-$ROOT_DIR/.flink-agents-env}" + local pyflink_jar="$FLINK_HOME/opt/flink-python-${FLINK_VERSION}.jar" + if [[ ! -f "$pyflink_jar" ]]; then + log_error "PyFlink JAR not found in Flink distribution: $pyflink_jar" + return 1 + fi + cp "$pyflink_jar" "$FLINK_HOME/lib/" + prepare_python_venv +} - log_info "Running tools/install.sh --non-interactive --install-flink --enable-pyflink" - FLINK_VERSION="$FLINK_VERSION" bash "$ROOT_DIR/tools/install.sh" \ - --non-interactive --install-flink --enable-pyflink +install_built_python_package() { + local wheel + wheel=$(find "$ROOT_DIR/python/dist" -maxdepth 1 -name '*.whl' | head -n 1) + if [[ -z "$wheel" ]]; then + log_error "Python wheel not found after build" + return 1 + fi - local install_dir="${INSTALL_DIR:-$HOME/.local/flink}" - export FLINK_HOME="${install_dir}/flink-${FLINK_VERSION}" + log_info "Installing current wheel and apache-flink==$FLINK_VERSION into $VENV_DIR" + "$VENV_DIR/bin/python" -m pip install --quiet \ + "$wheel" "apache-flink==$FLINK_VERSION" - if [[ ! -x "$FLINK_HOME/bin/flink" ]]; then - log_error "Flink installation not found at expected path: $FLINK_HOME" - exit 1 + if ! "$VENV_DIR/bin/python" -c 'import flink_agents, pyflink' >/dev/null 2>&1; then + log_error "Current Flink Agents wheel or PyFlink is not importable from: $VENV_DIR/bin/python" + return 1 fi - log_ok "Flink installed at: $FLINK_HOME" - # `flink run -py` shells out to a Python interpreter that must have - # pyflink importable. Activate the venv install.sh provisioned and - # point PYFLINK_CLIENT_EXECUTABLE at it. - if [[ ! -x "$VENV_DIR/bin/python" ]]; then - log_error "Expected Python venv not found at: $VENV_DIR" - exit 1 - fi # shellcheck disable=SC1091 source "$VENV_DIR/bin/activate" export PYFLINK_CLIENT_EXECUTABLE="$VENV_DIR/bin/python" - log_ok "Activated PyFlink venv: $VENV_DIR" + log_ok "Activated current Flink Agents and PyFlink venv: $VENV_DIR" } build_project() { @@ -180,6 +251,7 @@ build_project() { cd "$ROOT_DIR" SKIP_SPOTLESS_CHECK=true bash tools/build.sh ) + install_built_python_package log_ok "Build completed" } @@ -272,6 +344,77 @@ extract_job_id() { grep -Eo 'JobID [0-9a-f]{32}' "$1" | tail -n 1 | awk '{print $2}' } +# Starts a Java example in attached mode but leaves the Flink client in the +# background. Keeping the client alive preserves temporary resources created +# by the example until startup validation and cancellation have completed. +ATTACHED_SUBMISSION_PID="" +ATTACHED_SUBMISSION_JOB_ID="" +remove_submission_pid() { + local target_pid="$1" + local remaining_pids=() + local pid + + for pid in "${SUBMISSION_PIDS[@]:-}"; do + if [[ -n "$pid" && "$pid" != "$target_pid" ]]; then + remaining_pids+=("$pid") + fi + done + SUBMISSION_PIDS=("${remaining_pids[@]}") +} + +start_attached_java_submission() { + local class_name="$1" out="$2" + local elapsed=0 + + ATTACHED_SUBMISSION_PID="" + ATTACHED_SUBMISSION_JOB_ID="" + + timeout "$SUBMIT_TIMEOUT" "$FLINK_HOME/bin/flink" run \ + -c "$class_name" \ + "$EXAMPLES_JAR" >"$out" 2>&1 & + ATTACHED_SUBMISSION_PID=$! + SUBMISSION_PIDS+=("$ATTACHED_SUBMISSION_PID") + + while (( elapsed < SUBMIT_TIMEOUT )); do + ATTACHED_SUBMISSION_JOB_ID=$(extract_job_id "$out") || true + if [[ -n "$ATTACHED_SUBMISSION_JOB_ID" ]]; then + return 0 + fi + + if ! kill -0 "$ATTACHED_SUBMISSION_PID" 2>/dev/null; then + wait "$ATTACHED_SUBMISSION_PID" 2>/dev/null || true + remove_submission_pid "$ATTACHED_SUBMISSION_PID" + return 1 + fi + + sleep 1 + elapsed=$((elapsed + 1)) + done + + log_error "Timed out waiting for Java example job id" + kill "$ATTACHED_SUBMISSION_PID" 2>/dev/null || true + wait "$ATTACHED_SUBMISSION_PID" 2>/dev/null || true + remove_submission_pid "$ATTACHED_SUBMISSION_PID" + return 1 +} + +wait_for_attached_submission_client() { + local pid="$1" name="$2" + local elapsed=0 + + while kill -0 "$pid" 2>/dev/null && (( elapsed < 30 )); do + sleep 1 + elapsed=$((elapsed + 1)) + done + + if kill -0 "$pid" 2>/dev/null; then + log_warn "$name: attached Flink client did not exit after cleanup; stopping it" + kill "$pid" 2>/dev/null || true + fi + wait "$pid" 2>/dev/null || true + remove_submission_pid "$pid" +} + # --------------------------------------------------------------------------- # Check Flink logs for unexpected errors/exceptions after a job completes. # Inspired by Apache Flink's e2e test-scripts/common.sh approach. @@ -327,33 +470,208 @@ check_logs_for_errors() { return 0 } -wait_for_job_finish() { - local job_id="$1" name="$2" timeout_sec="${3:-$JOB_FINISH_TIMEOUT}" +extract_job_state() { + printf '%s\n' "$1" \ + | grep -Eo '"state"[[:space:]]*:[[:space:]]*"[^"]+"' \ + | sed -n '1{s/.*"\([^"]*\)"$/\1/p;}' +} + +# Sets JOB_STATE to the current state returned by Flink's REST API. +JOB_STATE="" +get_job_state() { + local job_id="$1" + local response + JOB_STATE="" + + response=$(curl -fsS "http://localhost:8081/jobs/$job_id"; 2>/dev/null) || return 1 + JOB_STATE=$(extract_job_state "$response") || return 1 + [[ -n "$JOB_STATE" ]] +} + +extract_available_slots() { + printf '%s\n' "$1" \ + | grep -Eo '"slots-available"[[:space:]]*:[[:space:]]*[0-9]+' \ + | sed -n '1{s/.*:[[:space:]]*//p;}' +} + +# Sets AVAILABLE_SLOTS to the number reported by Flink's overview endpoint. +AVAILABLE_SLOTS="" +get_available_slots() { + local response + AVAILABLE_SLOTS="" + + response=$(curl -fsS "http://localhost:8081/overview"; 2>/dev/null) || return 1 + AVAILABLE_SLOTS=$(extract_available_slots "$response") || return 1 + [[ -n "$AVAILABLE_SLOTS" ]] +} + +wait_for_cluster_slot_available() { + local name="$1" timeout_sec="${2:-$SLOT_RELEASE_TIMEOUT}" local elapsed=0 + while (( elapsed < timeout_sec )); do - local status - status=$("$FLINK_HOME/bin/flink" list -a 2>/dev/null | grep "$job_id" || true) - if echo "$status" | grep -q "FINISHED"; then - log_ok "$name reached FINISHED status" - sleep 2 # allow log flush - if check_logs_for_errors "$name"; then - log_ok "$name completed successfully (no unexpected errors in logs)" - else - log_warn "$name finished but with warnings in logs (non-fatal)" - fi + if get_available_slots && (( AVAILABLE_SLOTS > 0 )); then return 0 - elif echo "$status" | grep -q "FAILED\|CANCELED"; then - log_error "$name ended with unexpected status" - check_logs_for_errors "$name" || true - return 1 fi - sleep 5 - elapsed=$((elapsed + 5)) + sleep "$JOB_STATUS_POLL_INTERVAL" + elapsed=$((elapsed + JOB_STATUS_POLL_INTERVAL)) done - log_error "$name timed out after ${timeout_sec}s" + + log_warn "$name: cluster slot was not released within ${timeout_sec}s" return 1 } +# Python operators may take substantially longer than the bounded health check +# to tear down after cancellation. Restart the single-node standalone cluster +# instead of letting one slow cleanup starve every subsequently submitted job. +restart_cluster_after_cleanup() { + local name="$1" + + log_warn "$name: restarting the standalone cluster to recover its slot" + if ! "$FLINK_HOME/bin/stop-cluster.sh" >/dev/null 2>&1; then + log_warn "$name: stop-cluster.sh reported an error; attempting startup anyway" + fi + if ! "$FLINK_HOME/bin/start-cluster.sh" >/dev/null 2>&1; then + log_error "$name: failed to restart the standalone cluster" + return 1 + fi + if ! wait_for_cluster_slot_available "$name after cluster restart" 60; then + log_error "$name: cluster restart did not recover an available slot" + return 1 + fi + + log_ok "$name: standalone cluster restarted and slot recovered" +} + +ensure_cluster_slot_available() { + local name="$1" + + if wait_for_cluster_slot_available "$name"; then + return 0 + fi + restart_cluster_after_cleanup "$name" +} + +# A detached example is considered healthy when it finishes successfully or +# remains continuously RUNNING for a short period. This catches startup +# failures without waiting for slow inference or unbounded sources to finish. +wait_for_job_healthy() { + local job_id="$1" name="$2" timeout_sec="${3:-$JOB_STARTUP_TIMEOUT}" + local stable_sec="${4:-$JOB_STABLE_RUNNING_SECONDS}" + local elapsed=0 + local running_for=0 + local previous_state="" + + while (( elapsed < timeout_sec )); do + local status + if get_job_state "$job_id"; then + status="$JOB_STATE" + else + status="UNAVAILABLE" + fi + + if [[ "$status" != "$previous_state" ]]; then + log_info "$name state: $status" + fi + + case "$status" in + FINISHED) + log_ok "$name reached FINISHED status" + return 0 + ;; + RUNNING) + if [[ "$previous_state" == "RUNNING" ]]; then + running_for=$((running_for + JOB_STATUS_POLL_INTERVAL)) + else + running_for=0 + fi + if (( running_for >= stable_sec )); then + log_ok "$name remained RUNNING for ${stable_sec}s" + return 0 + fi + ;; + FAILING|FAILED|CANCELLING|CANCELED|SUSPENDED) + log_error "$name entered unexpected state: $status" + check_logs_for_errors "$name" || true + return 1 + ;; + *) + running_for=0 + ;; + esac + + previous_state="$status" + sleep "$JOB_STATUS_POLL_INTERVAL" + elapsed=$((elapsed + JOB_STATUS_POLL_INTERVAL)) + done + + log_error "$name did not become stably RUNNING or FINISHED within ${timeout_sec}s" + return 1 +} + +cancel_job_after_check() { + local job_id="$1" name="$2" + + if get_job_state "$job_id"; then + case "$JOB_STATE" in + FINISHED|FAILED|CANCELED|SUSPENDED) + ensure_cluster_slot_available "$name" + return $? + ;; + esac + fi + + log_info "$name: cancelling job after startup verification" + if timeout 30 "$FLINK_HOME/bin/flink" cancel "$job_id" >/dev/null 2>&1; then + if ensure_cluster_slot_available "$name"; then + log_ok "$name cancelled and cluster slot released" + return 0 + fi + return 1 + fi + + # The job may have reached a terminal state while the cancel request raced + # with completion. Treat that as successful cleanup. + if get_job_state "$job_id"; then + case "$JOB_STATE" in + FINISHED|FAILED|CANCELED|SUSPENDED) + ensure_cluster_slot_available "$name" + return $? + ;; + esac + fi + + # A fast FINISHED job can disappear from the active-job endpoint before + # the failed cancel request returns. In this sequential single-slot test, + # an available slot is sufficient evidence that cleanup has completed. + if ensure_cluster_slot_available "$name"; then + log_ok "$name no longer occupies the cluster slot" + return 0 + fi + + log_error "$name: failed to cancel job $job_id" + return 1 +} + +verify_submitted_job() { + local job_id="$1" name="$2" + local healthy=0 + local cleaned_up=0 + + if wait_for_job_healthy "$job_id" "$name"; then + healthy=1 + fi + if cancel_job_after_check "$job_id" "$name"; then + cleaned_up=1 + fi + + if (( healthy == 1 && cleaned_up == 1 )); then + record_result "$name" "PASS" + else + record_result "$name" "FAIL" + fi +} + submit_java_example() { local class_name="$1" local label="java:${class_name##*.}" @@ -362,10 +680,7 @@ submit_java_example() { local out out=$(mktemp) local rc=0 - timeout "$SUBMIT_TIMEOUT" "$FLINK_HOME/bin/flink" run \ - --detached \ - -c "$class_name" \ - "$EXAMPLES_JAR" >"$out" 2>&1 || rc=$? + start_attached_java_submission "$class_name" "$out" || rc=$? cat "$out" if (( rc != 0 )); then @@ -375,26 +690,19 @@ submit_java_example() { return 0 fi - local job_id - job_id=$(extract_job_id "$out") || true - rm -f "$out" + local job_id="$ATTACHED_SUBMISSION_JOB_ID" if [[ -z "$job_id" ]]; then log_error "$label: could not extract job id" record_result "$label" "FAIL" + rm -f "$out" return 0 fi SUBMITTED_JOB_IDS+=("$job_id") - # Submission succeeded — record PASS immediately log_ok "$label submitted successfully (JobID: $job_id)" - record_result "$label" "PASS" - - # Optionally wait and report final status (informational only, does not affect PASS/FAIL) - if wait_for_job_finish "$job_id" "$label" "$JOB_FINISH_TIMEOUT"; then - log_ok "$label job reached FINISHED status" - else - log_warn "$label job did not reach FINISHED (expected with lightweight CI model)" - fi + verify_submitted_job "$job_id" "$label" + wait_for_attached_submission_client "$ATTACHED_SUBMISSION_PID" "$label" + rm -f "$out" } submit_python_example() { @@ -433,16 +741,8 @@ submit_python_example() { fi SUBMITTED_JOB_IDS+=("$job_id") - # Submission succeeded — record PASS log_ok "$label submitted successfully (JobID: $job_id)" - record_result "$label" "PASS" - - # Informational: wait and report final status - if wait_for_job_finish "$job_id" "$label" "$JOB_FINISH_TIMEOUT"; then - log_ok "$label job reached FINISHED status" - else - log_warn "$label job did not reach FINISHED (expected with lightweight CI model)" - fi + verify_submitted_job "$job_id" "$label" } # RAG examples run end-to-end (not as detached jobs) via flink run -py. @@ -479,16 +779,27 @@ submit_python_rag_example() { discover_java_examples() { # Find all example classes by scanning the jar's manifest or known package # Convention: all classes directly under org.apache.flink.agents.examples that end with "Example" - local classes=() - while IFS= read -r class; do - classes+=("$class") - done < <(jar -tf "$EXAMPLES_JAR" \ + local jar_entries + if ! jar_entries=$(jar -tf "$EXAMPLES_JAR"); then + log_error "Failed to inspect examples JAR: $EXAMPLES_JAR" + return 1 + fi + + local discovered_classes + discovered_classes=$(printf '%s\n' "$jar_entries" \ | grep '^org/apache/flink/agents/examples/[^/]*Example\.class$' \ - | sed 's|/|.|g; s|\.class$||') + | sed 's|/|.|g; s|\.class$||' || true) + + local classes=() + if [[ -n "$discovered_classes" ]]; then + while IFS= read -r class; do + classes+=("$class") + done <<< "$discovered_classes" + fi if [[ ${#classes[@]} -eq 0 ]]; then log_error "No Java example classes found in $EXAMPLES_JAR" - exit 1 + return 1 fi log_info "Discovered ${#classes[@]} Java example(s): ${classes[*]}" printf '%s\n' "${classes[@]}" @@ -496,14 +807,29 @@ discover_java_examples() { discover_python_quickstart_examples() { local dir="$ROOT_DIR/python/flink_agents/examples/quickstart" + if [[ ! -d "$dir" ]]; then + log_error "Python quickstart examples directory not found: $dir" + return 1 + fi + + local discovered_scripts + if ! discovered_scripts=$( + find "$dir" -maxdepth 1 -name '*_example.py' -type f | sort + ); then + log_error "Failed to discover Python quickstart examples in $dir" + return 1 + fi + local scripts=() - while IFS= read -r f; do - scripts+=("$f") - done < <(find "$dir" -maxdepth 1 -name '*_example.py' -type f | sort) + if [[ -n "$discovered_scripts" ]]; then + while IFS= read -r f; do + scripts+=("$f") + done <<< "$discovered_scripts" + fi if [[ ${#scripts[@]} -eq 0 ]]; then log_error "No Python quickstart examples found in $dir" - exit 1 + return 1 fi log_info "Discovered ${#scripts[@]} Python quickstart example(s)" printf '%s\n' "${scripts[@]}" @@ -513,16 +839,27 @@ discover_python_rag_examples() { local dir="$ROOT_DIR/python/flink_agents/examples/rag" if [[ ! -d "$dir" ]]; then log_info "No RAG examples directory found, skipping" - return + return 0 fi + + local discovered_scripts + if ! discovered_scripts=$( + find "$dir" -maxdepth 1 -name '*_example.py' -type f | sort + ); then + log_error "Failed to discover Python RAG examples in $dir" + return 1 + fi + local scripts=() - while IFS= read -r f; do - scripts+=("$f") - done < <(find "$dir" -maxdepth 1 -name '*_example.py' -type f | sort) + if [[ -n "$discovered_scripts" ]]; then + while IFS= read -r f; do + scripts+=("$f") + done <<< "$discovered_scripts" + fi if [[ ${#scripts[@]} -eq 0 ]]; then log_info "No RAG examples found in $dir" - return + return 0 fi log_info "Discovered ${#scripts[@]} Python RAG example(s)" printf '%s\n' "${scripts[@]}" @@ -531,11 +868,14 @@ discover_python_rag_examples() { setup_rag_knowledge_base() { local setup_script="$ROOT_DIR/python/flink_agents/examples/rag/knowledge_base_setup.py" if [[ ! -f "$setup_script" ]]; then - log_warn "RAG knowledge_base_setup.py not found, skipping RAG setup" + log_error "RAG knowledge_base_setup.py not found: $setup_script" return 1 fi log_info "Setting up RAG knowledge base" - python "$setup_script" || { log_error "RAG knowledge base setup failed"; return 1; } + if ! python "$setup_script"; then + log_error "RAG knowledge base setup failed" + return 1 + fi log_ok "RAG knowledge base ready" } @@ -549,25 +889,53 @@ main() { # Auto-discover and submit Java examples log_section "Step 7: submit Java examples" + local java_examples + if ! java_examples=$(discover_java_examples); then + log_error "Java example discovery failed" + return 1 + fi + if [[ -z "$java_examples" ]]; then + log_error "Java example discovery returned no examples" + return 1 + fi while IFS= read -r class; do submit_java_example "$class" - done < <(discover_java_examples) + done <<< "$java_examples" # Auto-discover and submit Python quickstart examples log_section "Step 8: submit Python quickstart examples" + local python_quickstart_examples + if ! python_quickstart_examples=$(discover_python_quickstart_examples); then + log_error "Python quickstart example discovery failed" + return 1 + fi + if [[ -z "$python_quickstart_examples" ]]; then + log_error "Python quickstart example discovery returned no examples" + return 1 + fi while IFS= read -r script; do submit_python_example "$script" - done < <(discover_python_quickstart_examples) + done <<< "$python_quickstart_examples" # Auto-discover and run Python RAG examples (these run end-to-end, not as detached jobs) log_section "Step 9: run Python RAG examples" - if setup_rag_knowledge_base; then + local python_rag_examples + if ! python_rag_examples=$(discover_python_rag_examples); then + log_error "Python RAG example discovery failed" + return 1 + fi + if [[ -n "$python_rag_examples" ]]; then + if ! setup_rag_knowledge_base; then + log_error "Cannot run Python RAG examples because setup failed" + return 1 + fi while IFS= read -r script; do submit_python_rag_example "$script" - done < <(discover_python_rag_examples) - else - log_warn "Skipping RAG examples due to setup failure" + done <<< "$python_rag_examples" fi } -main "$@" +if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then + trap cleanup EXIT + main "$@" +fi diff --git a/examples/src/main/java/org/apache/flink/agents/examples/agents/ParallelChatAgent.java b/examples/src/main/java/org/apache/flink/agents/examples/agents/ParallelChatAgent.java index c2984e44..2286b6f9 100644 --- a/examples/src/main/java/org/apache/flink/agents/examples/agents/ParallelChatAgent.java +++ b/examples/src/main/java/org/apache/flink/agents/examples/agents/ParallelChatAgent.java @@ -94,7 +94,8 @@ public class ParallelChatAgent extends Agent { return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.OLLAMA_SETUP) .addInitialArgument("connection", "ollamaChatModelConnection") .addInitialArgument("model", OLLAMA_MODEL) - .addInitialArgument("extract_reasoning", true) + .addInitialArgument("think", false) + .addInitialArgument("extract_reasoning", false) .build(); } diff --git a/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py b/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py index 6d98cb2a..f2e20652 100644 --- a/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py +++ b/python/flink_agents/examples/quickstart/agents/parallel_chat_agent.py @@ -114,7 +114,8 @@ class ParallelChatAgent(Agent): clazz=ResourceName.ChatModel.OLLAMA_SETUP, connection="ollama_server", model=OLLAMA_MODEL, - extract_reasoning=True, + think=False, + extract_reasoning=False, ) @action(InputEvent.EVENT_TYPE) @@ -155,7 +156,11 @@ class ParallelChatAgent(Agent): response_event = ChatResponseEvent.from_event(event) parsed = response_event.response.extra_args[STRUCTURED_OUTPUT] if isinstance(parsed, dict): - parsed = SummaryResponse(**parsed) if "summary" in parsed else AspectResponse(**parsed) + parsed = ( + SummaryResponse(**parsed) + if "summary" in parsed + else AspectResponse(**parsed) + ) if isinstance(parsed, SummaryResponse): ctx.send_event( _build_output_event( diff --git a/python/flink_agents/examples/rag/agents/__init__.py b/python/flink_agents/examples/rag/agents/__init__.py new file mode 100644 index 00000000..65b48d4d --- /dev/null +++ b/python/flink_agents/examples/rag/agents/__init__.py @@ -0,0 +1,17 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ diff --git a/python/flink_agents/examples/rag/rag_agent_example.py b/python/flink_agents/examples/rag/agents/rag_agent.py similarity index 74% copy from python/flink_agents/examples/rag/rag_agent_example.py copy to python/flink_agents/examples/rag/agents/rag_agent.py index 00e688b6..461080d1 100644 --- a/python/flink_agents/examples/rag/rag_agent_example.py +++ b/python/flink_agents/examples/rag/agents/rag_agent.py @@ -16,8 +16,8 @@ # limitations under the License. ################################################################################ import os - -from pyflink.datastream import StreamExecutionEnvironment +import tempfile +from pathlib import Path from flink_agents.api.agents.agent import Agent from flink_agents.api.chat_message import ChatMessage, MessageRole @@ -35,7 +35,6 @@ from flink_agents.api.events.context_retrieval_event import ( ) from flink_agents.api.events.event import Event, InputEvent, OutputEvent from flink_agents.api.events.event_type import EventType -from flink_agents.api.execution_environment import AgentsExecutionEnvironment from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.resource import ( ResourceDescriptor, @@ -43,10 +42,13 @@ from flink_agents.api.resource import ( ResourceType, ) from flink_agents.api.runner_context import RunnerContext -from flink_agents.examples.rag.knowledge_base_setup import populate_knowledge_base OLLAMA_CHAT_MODEL = os.environ.get("OLLAMA_CHAT_MODEL", "qwen3:8b") OLLAMA_EMBEDDING_MODEL = os.environ.get("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text") +CHROMA_PERSIST_DIRECTORY = os.environ.get( + "CHROMA_PERSIST_DIRECTORY", + str(Path(tempfile.gettempdir()) / "flink-agents-rag-chroma"), +) class MyRAGAgent(Agent): @@ -94,6 +96,7 @@ Please provide a helpful answer based on the context provided.""" clazz=ResourceName.VectorStore.CHROMA_VECTOR_STORE, embedding_model="text_embedder", collection="example_knowledge_base", + persist_directory=CHROMA_PERSIST_DIRECTORY, ) @chat_model_setup @@ -104,6 +107,7 @@ Please provide a helpful answer based on the context provided.""" clazz=ResourceName.ChatModel.OLLAMA_SETUP, connection="ollama_chat_connection", model=OLLAMA_CHAT_MODEL, + think=False, ) @action(EventType.InputEvent) @@ -121,20 +125,16 @@ Please provide a helpful answer based on the context provided.""" @action(EventType.ContextRetrievalResponseEvent) @staticmethod - def process_retrieved_context( - event: Event, ctx: RunnerContext - ) -> None: + def process_retrieved_context(event: Event, ctx: RunnerContext) -> None: """Process retrieved context and create enhanced chat request.""" retrieval_event = ContextRetrievalResponseEvent.from_event(event) user_query = retrieval_event.query retrieved_docs = retrieval_event.documents - # Create context from retrieved documents context_text = "\n\n".join( [f"{i + 1}. {doc.content}" for i, doc in enumerate(retrieved_docs)] ) - # Get prompt resource and format it prompt_resource = ctx.get_resource( "context_enhanced_prompt", ResourceType.PROMPT ) @@ -142,7 +142,6 @@ Please provide a helpful answer based on the context provided.""" context=context_text, user_query=user_query ) - # Send chat request with enhanced prompt ctx.send_event( ChatRequestEvent( model="chat_model", @@ -157,48 +156,3 @@ Please provide a helpful answer based on the context provided.""" chat_response = ChatResponseEvent.from_event(event) if chat_response.response and chat_response.response.content: ctx.send_event(OutputEvent(output=chat_response.response.content)) - - -if __name__ == "__main__": - print("Starting RAG Example Agent...") - - # Populate vector store with sample documents - populate_knowledge_base() - - agent = MyRAGAgent() - - # Set up the Flink streaming environment and the Agents execution environment. - env = StreamExecutionEnvironment.get_execution_environment() - agents_env = AgentsExecutionEnvironment.get_execution_environment(env) - - # Setup Ollama embedding and chat model connections - agents_env.add_resource( - "ollama_embedding_connection", - ResourceType.EMBEDDING_MODEL_CONNECTION, - ResourceDescriptor(clazz=ResourceName.EmbeddingModel.OLLAMA_CONNECTION), - ) - agents_env.add_resource( - "ollama_chat_connection", - ResourceType.EMBEDDING_MODEL, - ResourceDescriptor(clazz=ResourceName.ChatModel.OLLAMA_CONNECTION), - ) - - # A small stream of example queries, keyed by the query text. - query_stream = env.from_collection( - [ - "What is Apache Flink?", - "What is Apache Flink Agents?", - "What is Python?", - ], - ) - - # Use the RAG agent to answer each query and print the responses to stdout. - response_stream = ( - agents_env.from_datastream(input=query_stream, key_selector=lambda x: x) - .apply(agent) - .to_datastream() - ) - response_stream.print() - - # Execute the Flink pipeline. - agents_env.execute("RAG Agent Example Job") diff --git a/python/flink_agents/examples/rag/knowledge_base_setup.py b/python/flink_agents/examples/rag/knowledge_base_setup.py index de3f6f05..ace827ec 100644 --- a/python/flink_agents/examples/rag/knowledge_base_setup.py +++ b/python/flink_agents/examples/rag/knowledge_base_setup.py @@ -17,6 +17,8 @@ ################################################################################# import os +import tempfile +from pathlib import Path import chromadb @@ -25,6 +27,10 @@ from flink_agents.integrations.embedding_models.local.ollama_embedding_model imp ) OLLAMA_EMBEDDING_MODEL = os.environ.get("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text") +CHROMA_PERSIST_DIRECTORY = os.environ.get( + "CHROMA_PERSIST_DIRECTORY", + str(Path(tempfile.gettempdir()) / "flink-agents-rag-chroma"), +) """Utility for populating ChromaDB with sample knowledge documents for RAG examples.""" @@ -35,7 +41,10 @@ def populate_knowledge_base() -> None: # Create connections directly embedding_connection = OllamaEmbeddingModelConnection() - chroma_client = chromadb.EphemeralClient() + # The Flink client and TaskManager run in separate processes, so an + # in-memory Chroma client cannot share the populated collection with the + # agent. Persist it on the local filesystem used by the standalone cluster. + chroma_client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIRECTORY) # Get collection (create if doesn't exist) collection_name = "example_knowledge_base" @@ -76,7 +85,7 @@ def populate_knowledge_base() -> None: } # Add documents to ChromaDB - collection.add(**test_data) + collection.upsert(**test_data) print( f"Knowledge base setup complete! Added {len(documents)} documents to ChromaDB." diff --git a/python/flink_agents/examples/rag/rag_agent_example.py b/python/flink_agents/examples/rag/rag_agent_example.py index 00e688b6..2a90d669 100644 --- a/python/flink_agents/examples/rag/rag_agent_example.py +++ b/python/flink_agents/examples/rag/rag_agent_example.py @@ -15,150 +15,17 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################ -import os - from pyflink.datastream import StreamExecutionEnvironment -from flink_agents.api.agents.agent import Agent -from flink_agents.api.chat_message import ChatMessage, MessageRole -from flink_agents.api.decorators import ( - action, - chat_model_setup, - embedding_model_setup, - prompt, - vector_store, -) -from flink_agents.api.events.chat_event import ChatRequestEvent, ChatResponseEvent -from flink_agents.api.events.context_retrieval_event import ( - ContextRetrievalRequestEvent, - ContextRetrievalResponseEvent, -) -from flink_agents.api.events.event import Event, InputEvent, OutputEvent -from flink_agents.api.events.event_type import EventType from flink_agents.api.execution_environment import AgentsExecutionEnvironment -from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.resource import ( ResourceDescriptor, ResourceName, ResourceType, ) -from flink_agents.api.runner_context import RunnerContext +from flink_agents.examples.rag.agents.rag_agent import MyRAGAgent from flink_agents.examples.rag.knowledge_base_setup import populate_knowledge_base -OLLAMA_CHAT_MODEL = os.environ.get("OLLAMA_CHAT_MODEL", "qwen3:8b") -OLLAMA_EMBEDDING_MODEL = os.environ.get("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text") - - -class MyRAGAgent(Agent): - """Example RAG agent demonstrating context retrieval. - - This RAG agent shows how to: - 1. Receive a user query - 2. Retrieve relevant context from a vector store using semantic search - 3. Augment the user query with retrieved context - 4. Generate enhanced responses using the chat model - - This is a basic example demonstrating RAG workflow with Ollama and ChromaDB. - """ - - @prompt - @staticmethod - def context_enhanced_prompt() -> Prompt: - """Prompt template for enhancing user queries with retrieved context.""" - template = """Based on the following context, please answer the user's question. - -Context: -{context} - -User Question: -{user_query} - -Please provide a helpful answer based on the context provided.""" - return Prompt.from_text(template) - - @embedding_model_setup - @staticmethod - def text_embedder() -> ResourceDescriptor: - """Embedding model setup for generating text embeddings.""" - return ResourceDescriptor( - clazz=ResourceName.EmbeddingModel.OLLAMA_SETUP, - connection="ollama_embedding_connection", - model=OLLAMA_EMBEDDING_MODEL, - ) - - @vector_store - @staticmethod - def knowledge_base() -> ResourceDescriptor: - """Vector store setup for knowledge base.""" - return ResourceDescriptor( - clazz=ResourceName.VectorStore.CHROMA_VECTOR_STORE, - embedding_model="text_embedder", - collection="example_knowledge_base", - ) - - @chat_model_setup - @staticmethod - def chat_model() -> ResourceDescriptor: - """Chat model setup for generating responses.""" - return ResourceDescriptor( - clazz=ResourceName.ChatModel.OLLAMA_SETUP, - connection="ollama_chat_connection", - model=OLLAMA_CHAT_MODEL, - ) - - @action(EventType.InputEvent) - @staticmethod - def process_input(event: Event, ctx: RunnerContext) -> None: - """Process user input and retrieve relevant context.""" - user_query = str(InputEvent.from_event(event).input) - ctx.send_event( - ContextRetrievalRequestEvent( - query=user_query, - vector_store="knowledge_base", - max_results=3, - ) - ) - - @action(EventType.ContextRetrievalResponseEvent) - @staticmethod - def process_retrieved_context( - event: Event, ctx: RunnerContext - ) -> None: - """Process retrieved context and create enhanced chat request.""" - retrieval_event = ContextRetrievalResponseEvent.from_event(event) - user_query = retrieval_event.query - retrieved_docs = retrieval_event.documents - - # Create context from retrieved documents - context_text = "\n\n".join( - [f"{i + 1}. {doc.content}" for i, doc in enumerate(retrieved_docs)] - ) - - # Get prompt resource and format it - prompt_resource = ctx.get_resource( - "context_enhanced_prompt", ResourceType.PROMPT - ) - enhanced_prompt = prompt_resource.format_string( - context=context_text, user_query=user_query - ) - - # Send chat request with enhanced prompt - ctx.send_event( - ChatRequestEvent( - model="chat_model", - messages=[ChatMessage(role=MessageRole.USER, content=enhanced_prompt)], - ) - ) - - @action(EventType.ChatResponseEvent) - @staticmethod - def process_chat_response(event: Event, ctx: RunnerContext) -> None: - """Process chat model response and generate output.""" - chat_response = ChatResponseEvent.from_event(event) - if chat_response.response and chat_response.response.content: - ctx.send_event(OutputEvent(output=chat_response.response.content)) - - if __name__ == "__main__": print("Starting RAG Example Agent...") @@ -179,8 +46,11 @@ if __name__ == "__main__": ) agents_env.add_resource( "ollama_chat_connection", - ResourceType.EMBEDDING_MODEL, - ResourceDescriptor(clazz=ResourceName.ChatModel.OLLAMA_CONNECTION), + ResourceType.CHAT_MODEL_CONNECTION, + ResourceDescriptor( + clazz=ResourceName.ChatModel.OLLAMA_CONNECTION, + request_timeout=240.0, + ), ) # A small stream of example queries, keyed by the query text. diff --git a/python/flink_agents/examples/rag/tests/test_rag_agent.py b/python/flink_agents/examples/rag/tests/test_rag_agent.py new file mode 100644 index 00000000..db75f3af --- /dev/null +++ b/python/flink_agents/examples/rag/tests/test_rag_agent.py @@ -0,0 +1,46 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ + +from flink_agents.api.function import PythonFunction +from flink_agents.examples.rag.agents.rag_agent import ( + CHROMA_PERSIST_DIRECTORY as AGENT_CHROMA_DIRECTORY, +) +from flink_agents.examples.rag.agents.rag_agent import MyRAGAgent +from flink_agents.examples.rag.knowledge_base_setup import ( + CHROMA_PERSIST_DIRECTORY as SETUP_CHROMA_DIRECTORY, +) + + +def test_rag_agent_actions_resolve_from_an_importable_module() -> None: + actions = [ + MyRAGAgent.process_input, + MyRAGAgent.process_retrieved_context, + MyRAGAgent.process_chat_response, + ] + + for action in actions: + descriptor = PythonFunction.from_callable(action) + assert descriptor.module == "flink_agents.examples.rag.agents.rag_agent" + assert descriptor.as_callable() is action + + +def test_rag_vector_store_uses_cross_process_persistence() -> None: + vector_store = MyRAGAgent.knowledge_base() + + assert AGENT_CHROMA_DIRECTORY == SETUP_CHROMA_DIRECTORY + assert vector_store.arguments["persist_directory"] == SETUP_CHROMA_DIRECTORY diff --git a/tools/test/unit/verify_example_job.bats b/tools/test/unit/verify_example_job.bats new file mode 100644 index 00000000..2ed65a77 --- /dev/null +++ b/tools/test/unit/verify_example_job.bats @@ -0,0 +1,466 @@ +#!/usr/bin/env bats +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +setup() { + source "${BATS_TEST_DIRNAME}/../../../e2e-test/test-scripts/test_submit_examples_to_flink.sh" + JOB_STATUS_POLL_INTERVAL=1 + TEST_STATE_INDEX=0 + TEST_STATES=() + + sleep() { + : + } + + get_job_state() { + local last_index=$((${#TEST_STATES[@]} - 1)) + local index="$TEST_STATE_INDEX" + if (( index > last_index )); then + index="$last_index" + fi + JOB_STATE="${TEST_STATES[$index]}" + TEST_STATE_INDEX=$((TEST_STATE_INDEX + 1)) + return 0 + } +} + +set_job_states() { + TEST_STATES=("$@") + TEST_STATE_INDEX=0 +} + +stub_main_dependencies() { + install_flink() { :; } + build_project() { :; } + locate_examples_jar() { :; } + stage_dist_jars() { :; } + start_ollama() { :; } + start_cluster() { :; } + discover_java_examples() { printf 'org.example.JavaExample\n'; } + discover_python_quickstart_examples() { printf '/tmp/python_example.py\n'; } + discover_python_rag_examples() { printf '/tmp/rag_example.py\n'; } + setup_rag_knowledge_base() { :; } + submit_java_example() { :; } + submit_python_example() { :; } + submit_python_rag_example() { :; } +} + +create_fake_executable() { + local path="$1" + mkdir -p "$(dirname "$path")" + printf '#!/usr/bin/env bash\nexit 0\n' > "$path" + chmod +x "$path" +} + +@test "job state parser extracts the current state" { + run extract_job_state '{"jid":"job-id","name":"example","state":"RUNNING"}' + + [ "$status" -eq 0 ] + [ "$output" = "RUNNING" ] +} + +@test "job state parser rejects a response without state" { + run extract_job_state '{"errors":["job not found"]}' + + [ "$status" -ne 0 ] + [ -z "$output" ] +} + +@test "available slot parser extracts the current count" { + run extract_available_slots '{"taskmanagers":1,"slots-total":1,"slots-available":1}' + + [ "$status" -eq 0 ] + [ "$output" = "1" ] +} + +@test "slot release check waits until a slot becomes available" { + local slot_index=0 + get_available_slots() { + local slots=(0 0 1) + AVAILABLE_SLOTS="${slots[$slot_index]}" + if (( slot_index < 2 )); then + slot_index=$((slot_index + 1)) + fi + return 0 + } + + run wait_for_cluster_slot_available "example" 4 + + [ "$status" -eq 0 ] +} + +@test "job health check accepts FINISHED immediately" { + set_job_states FINISHED + + run wait_for_job_healthy "job-id" "example" 5 2 + + [ "$status" -eq 0 ] + [[ "$output" == *"reached FINISHED status"* ]] +} + +@test "job health check accepts a continuously RUNNING job" { + set_job_states CREATED RUNNING RUNNING RUNNING + + run wait_for_job_healthy "job-id" "example" 5 2 + + [ "$status" -eq 0 ] + [[ "$output" == *"remained RUNNING for 2s"* ]] +} + +@test "job health check rejects a failing job" { + set_job_states CREATED RUNNING FAILING + check_logs_for_errors() { + return 0 + } + + run wait_for_job_healthy "job-id" "example" 5 2 + + [ "$status" -ne 0 ] + [[ "$output" == *"entered unexpected state: FAILING"* ]] +} + +@test "job health check resets the stability period after a restart" { + set_job_states RUNNING RESTARTING RUNNING RUNNING + + run wait_for_job_healthy "job-id" "example" 4 2 + + [ "$status" -ne 0 ] + [[ "$output" == *"did not become stably RUNNING or FINISHED"* ]] +} + +@test "submitted job is marked failed when health verification fails" { + RESULT_NAMES=() + RESULT_STATES=() + wait_for_job_healthy() { + return 1 + } + cancel_job_after_check() { + return 0 + } + + verify_submitted_job "job-id" "example" + + [ "${RESULT_NAMES[0]}" = "example" ] + [ "${RESULT_STATES[0]}" = "FAIL" ] +} + +@test "submitted job passes only after health verification and cleanup" { + RESULT_NAMES=() + RESULT_STATES=() + CANCEL_CALLED=0 + wait_for_job_healthy() { + return 0 + } + cancel_job_after_check() { + CANCEL_CALLED=1 + return 0 + } + + verify_submitted_job "job-id" "example" + + [ "$CANCEL_CALLED" -eq 1 ] + [ "${RESULT_NAMES[0]}" = "example" ] + [ "${RESULT_STATES[0]}" = "PASS" ] +} + +@test "submitted job is marked failed when cleanup fails" { + RESULT_NAMES=() + RESULT_STATES=() + wait_for_job_healthy() { + return 0 + } + cancel_job_after_check() { + return 1 + } + + verify_submitted_job "job-id" "example" + + [ "${RESULT_STATES[0]}" = "FAIL" ] +} + +@test "cancel race succeeds when a finished job has already released its slot" { + FLINK_HOME="$BATS_TEST_TMPDIR/flink" + local state_calls=0 + get_job_state() { + state_calls=$((state_calls + 1)) + if (( state_calls == 1 )); then + JOB_STATE="RUNNING" + return 0 + fi + return 1 + } + timeout() { + return 1 + } + ensure_cluster_slot_available() { + return 0 + } + + run cancel_job_after_check "job-id" "example" + + [ "$status" -eq 0 ] + [[ "$output" == *"no longer occupies the cluster slot"* ]] +} + +@test "slow slot cleanup restarts the standalone cluster" { + local restarted=0 + wait_for_cluster_slot_available() { + return 1 + } + restart_cluster_after_cleanup() { + restarted=1 + return 0 + } + + ensure_cluster_slot_available "example" + + [ "$restarted" -eq 1 ] +} + +@test "cleanup fails when neither cancellation nor cluster recovery succeeds" { + FLINK_HOME="$BATS_TEST_TMPDIR/flink" + get_job_state() { + JOB_STATE="RUNNING" + return 0 + } + timeout() { + return 1 + } + ensure_cluster_slot_available() { + return 1 + } + + run cancel_job_after_check "job-id" "example" + + [ "$status" -ne 0 ] + [[ "$output" == *"failed to cancel job job-id"* ]] +} + +@test "CI chat model aliases cover every hardcoded quickstart model" { + [[ " ${OLLAMA_CHAT_MODEL_ALIASES[*]} " == *" qwen3:1.7b "* ]] + [[ " ${OLLAMA_CHAT_MODEL_ALIASES[*]} " == *" qwen3:8b "* ]] + [[ " ${OLLAMA_CHAT_MODEL_ALIASES[*]} " == *" qwen3.5:9b "* ]] +} + +@test "Java submission keeps the attached client alive during validation" { + unset -f sleep + + local fake_flink_home="$BATS_TEST_TMPDIR/flink" + local client_marker="$BATS_TEST_TMPDIR/client-running" + local client_release="$BATS_TEST_TMPDIR/client-release" + mkdir -p "$fake_flink_home/bin" + printf '%s\n' \ + '#!/usr/bin/env bash' \ + 'touch "$ATTACHED_CLIENT_MARKER"' \ + 'printf "Job has been submitted with JobID 0123456789abcdef0123456789abcdef\n"' \ + 'while [[ ! -f "$ATTACHED_CLIENT_RELEASE" ]]; do sleep 0.1; done' \ + > "$fake_flink_home/bin/flink" + chmod +x "$fake_flink_home/bin/flink" + + FLINK_HOME="$fake_flink_home" + EXAMPLES_JAR="$BATS_TEST_TMPDIR/examples.jar" + ATTACHED_CLIENT_MARKER="$client_marker" + ATTACHED_CLIENT_RELEASE="$client_release" + export ATTACHED_CLIENT_MARKER ATTACHED_CLIENT_RELEASE + SUBMIT_TIMEOUT=5 + RESULT_NAMES=() + RESULT_STATES=() + SUBMITTED_JOB_IDS=() + SUBMISSION_PIDS=() + + timeout() { + shift + "$@" + } + wait_for_job_healthy() { + [[ -f "$client_marker" ]] + } + cancel_job_after_check() { + touch "$client_release" + } + + submit_java_example "org.example.JavaExample" + + [ "${RESULT_NAMES[0]}" = "java:JavaExample" ] + [ "${RESULT_STATES[0]}" = "PASS" ] + [ "${SUBMITTED_JOB_IDS[0]}" = "0123456789abcdef0123456789abcdef" ] + [ "${#SUBMISSION_PIDS[@]}" -eq 0 ] +} + +@test "main propagates Java example discovery failure" { + stub_main_dependencies + discover_java_examples() { + return 1 + } + + run main + + [ "$status" -ne 0 ] + [[ "$output" == *"Java example discovery failed"* ]] +} + +@test "main propagates Python quickstart discovery failure" { + stub_main_dependencies + discover_python_quickstart_examples() { + return 1 + } + + run main + + [ "$status" -ne 0 ] + [[ "$output" == *"Python quickstart example discovery failed"* ]] +} + +@test "main propagates RAG example discovery failure" { + stub_main_dependencies + discover_python_rag_examples() { + return 1 + } + + run main + + [ "$status" -ne 0 ] + [[ "$output" == *"Python RAG example discovery failed"* ]] +} + +@test "main propagates RAG knowledge base setup failure" { + stub_main_dependencies + setup_rag_knowledge_base() { + return 1 + } + + run main + + [ "$status" -ne 0 ] + [[ "$output" == *"Cannot run Python RAG examples because setup failed"* ]] +} + +@test "main skips RAG setup when no RAG examples exist" { + stub_main_dependencies + discover_python_rag_examples() { + return 0 + } + setup_rag_knowledge_base() { + return 1 + } + + run main + + [ "$status" -eq 0 ] +} + +@test "Flink bootstrap sources installer helpers without installing released Agents" { + ROOT_DIR="$BATS_TEST_TMPDIR/root" + BOOTSTRAP_CALLS="$BATS_TEST_TMPDIR/bootstrap-calls" + export BOOTSTRAP_CALLS + mkdir -p "$ROOT_DIR/tools" + printf '%s\n' \ + 'plan_flink() {' \ + ' printf "plan:%s:%s:%s\\n" "$FLINK_VERSION" "$INSTALL_FLINK" "$INSTALL_DIR" >> "$BOOTSTRAP_CALLS"' \ + '}' \ + 'install_flink_if_needed() {' \ + ' printf "install-flink\\n" >> "$BOOTSTRAP_CALLS"' \ + '}' \ + 'install_flink_agents_jar() {' \ + ' printf "install-agents\\n" >> "$BOOTSTRAP_CALLS"' \ + ' return 1' \ + '}' \ + > "$ROOT_DIR/tools/install.sh" + + install_flink_distribution "$BATS_TEST_TMPDIR/install" + + [ "$(sed -n '1p' "$BOOTSTRAP_CALLS")" = \ + "plan:$FLINK_VERSION:Yes:$BATS_TEST_TMPDIR/install" ] + [ "$(sed -n '2p' "$BOOTSTRAP_CALLS")" = "install-flink" ] + [ "$(wc -l < "$BOOTSTRAP_CALLS" | tr -d ' ')" -eq 2 ] +} + +@test "reused FLINK_HOME prepares PyFlink without installing released Agents artifacts" { + local fake_flink_home="$BATS_TEST_TMPDIR/flink" + local fake_venv="$BATS_TEST_TMPDIR/venv" + create_fake_executable "$fake_flink_home/bin/flink" + mkdir -p "$fake_flink_home/lib" "$fake_flink_home/opt" + : > "$fake_flink_home/opt/flink-python-${FLINK_VERSION}.jar" + + FLINK_HOME="$fake_flink_home" + VENV_DIR="$fake_venv" + PREPARE_VENV_CALLED=0 + INSTALL_DISTRIBUTION_CALLED=0 + prepare_python_venv() { + PREPARE_VENV_CALLED=1 + } + install_flink_distribution() { + INSTALL_DISTRIBUTION_CALLED=1 + return 1 + } + + install_flink + + [ "$INSTALL_DISTRIBUTION_CALLED" -eq 0 ] + [ "$PREPARE_VENV_CALLED" -eq 1 ] + [ -f "$fake_flink_home/lib/flink-python-${FLINK_VERSION}.jar" ] +} + +@test "fresh setup invokes only the Flink distribution installer" { + local install_dir="$BATS_TEST_TMPDIR/install" + local fake_flink_home="$install_dir/flink-$FLINK_VERSION" + local fake_venv="$BATS_TEST_TMPDIR/venv" + + unset FLINK_HOME + INSTALL_DIR="$install_dir" + VENV_DIR="$fake_venv" + INSTALL_DISTRIBUTION_CALLED=0 + prepare_python_venv() { + : + } + install_flink_distribution() { + INSTALL_DISTRIBUTION_CALLED=1 + [ "$1" = "$install_dir" ] + create_fake_executable "$fake_flink_home/bin/flink" + mkdir -p "$fake_flink_home/lib" "$fake_flink_home/opt" + : > "$fake_flink_home/opt/flink-python-${FLINK_VERSION}.jar" + } + + install_flink + + [ "$FLINK_HOME" = "$fake_flink_home" ] + [ "$INSTALL_DISTRIBUTION_CALLED" -eq 1 ] + [ -f "$fake_flink_home/lib/flink-python-${FLINK_VERSION}.jar" ] +} + +@test "built wheel and matching PyFlink are installed into the E2E venv" { + ROOT_DIR="$BATS_TEST_TMPDIR/root" + VENV_DIR="$BATS_TEST_TMPDIR/venv" + PYTHON_CALLS="$BATS_TEST_TMPDIR/python-calls" + export PYTHON_CALLS + mkdir -p "$ROOT_DIR/python/dist" "$VENV_DIR/bin" + : > "$ROOT_DIR/python/dist/flink_agents-0.3.0-py3-none-any.whl" + printf '%s\n' \ + '#!/usr/bin/env bash' \ + 'printf "%s\\n" "$*" >> "$PYTHON_CALLS"' \ + 'exit 0' \ + > "$VENV_DIR/bin/python" + chmod +x "$VENV_DIR/bin/python" + printf 'TEST_VENV_ACTIVATED=1\n' > "$VENV_DIR/bin/activate" + + install_built_python_package + + [ "$TEST_VENV_ACTIVATED" -eq 1 ] + [ "$PYFLINK_CLIENT_EXECUTABLE" = "$VENV_DIR/bin/python" ] + [[ "$(cat "$PYTHON_CALLS")" == *"-m pip install --quiet"* ]] + [[ "$(cat "$PYTHON_CALLS")" == *"apache-flink==$FLINK_VERSION"* ]] +}
