From 207c871375e3e91154ea04560ad73bfccb235d94 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Tue, 9 Sep 2025 12:43:40 -0700
Subject: [PATCH] fix shell quoting

---
 .../actions/run-and-record-tests/action.yml   | 10 +--
 scripts/get_setup_env.py                      | 71 +++++++++++++++++++
 scripts/integration-tests.sh                  | 41 ++++++-----
 3 files changed, 101 insertions(+), 21 deletions(-)
 create mode 100644 scripts/get_setup_env.py

diff --git a/.github/actions/run-and-record-tests/action.yml b/.github/actions/run-and-record-tests/action.yml
index 7af1a0040..a3eb31d9f 100644
--- a/.github/actions/run-and-record-tests/action.yml
+++ b/.github/actions/run-and-record-tests/action.yml
@@ -38,20 +38,20 @@ runs:
     - name: Run Integration Tests
       shell: bash
       run: |
-        SCRIPT_ARGS="--stack-config '${{ inputs.stack-config }}' --inference-mode '${{ inputs.inference-mode }}'"
+        SCRIPT_ARGS="--stack-config ${{ inputs.stack-config }} --inference-mode ${{ inputs.inference-mode }}"
 
         # Add optional arguments only if they are provided
         if [ -n '${{ inputs.setup }}' ]; then
-          SCRIPT_ARGS="$SCRIPT_ARGS --setup '${{ inputs.setup }}'"
+          SCRIPT_ARGS="$SCRIPT_ARGS --setup ${{ inputs.setup }}"
         fi
         if [ -n '${{ inputs.suite }}' ]; then
-          SCRIPT_ARGS="$SCRIPT_ARGS --suite '${{ inputs.suite }}'"
+          SCRIPT_ARGS="$SCRIPT_ARGS --suite ${{ inputs.suite }}"
         fi
         if [ -n '${{ inputs.subdirs }}' ]; then
-          SCRIPT_ARGS="$SCRIPT_ARGS --subdirs '${{ inputs.subdirs }}'"
+          SCRIPT_ARGS="$SCRIPT_ARGS --subdirs ${{ inputs.subdirs }}"
         fi
         if [ -n '${{ inputs.pattern }}' ]; then
-          SCRIPT_ARGS="$SCRIPT_ARGS --pattern '${{ inputs.pattern }}'"
+          SCRIPT_ARGS="$SCRIPT_ARGS --pattern ${{ inputs.pattern }}"
         fi
 
         uv run --no-sync ./scripts/integration-tests.sh $SCRIPT_ARGS | tee pytest-${{ inputs.inference-mode }}.log
diff --git a/scripts/get_setup_env.py b/scripts/get_setup_env.py
new file mode 100644
index 000000000..fad601e76
--- /dev/null
+++ b/scripts/get_setup_env.py
@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+"""
+Small helper script to extract environment variables from a test setup.
+Used by integration-tests.sh to set environment variables before starting the server.
+"""
+
+import argparse
+import sys
+
+from tests.integration.suites import SETUP_DEFINITIONS, SUITE_DEFINITIONS
+
+
+def get_setup_env_vars(setup_name, suite_name=None):
+    """
+    Get environment variables for a setup, with optional suite default fallback.
+
+    Args:
+        setup_name: Name of the setup (e.g., 'ollama', 'gpt')
+        suite_name: Optional suite name to get default setup if setup_name is None
+
+    Returns:
+        Dictionary of environment variables
+    """
+    # If no setup specified, try to get default from suite
+    if not setup_name and suite_name:
+        suite = SUITE_DEFINITIONS.get(suite_name)
+        if suite and suite.default_setup:
+            setup_name = suite.default_setup
+
+    if not setup_name:
+        return {}
+
+    setup = SETUP_DEFINITIONS.get(setup_name)
+    if not setup:
+        print(
+            f"Error: Unknown setup '{setup_name}'. Available: {', '.join(sorted(SETUP_DEFINITIONS.keys()))}",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    return setup.env
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Extract environment variables from a test setup")
+    parser.add_argument("--setup", help="Setup name (e.g., ollama, gpt)")
+    parser.add_argument("--suite", help="Suite name to get default setup from if --setup not provided")
+    parser.add_argument("--format", choices=["bash", "json"], default="bash", help="Output format (default: bash)")
+
+    args = parser.parse_args()
+
+    env_vars = get_setup_env_vars(args.setup, args.suite)
+
+    if args.format == "bash":
+        # Output as bash export statements
+        for key, value in env_vars.items():
+            print(f"export {key}='{value}'")
+    elif args.format == "json":
+        import json
+
+        print(json.dumps(env_vars))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/integration-tests.sh b/scripts/integration-tests.sh
index aa98a351a..0056bb842 100755
--- a/scripts/integration-tests.sh
+++ b/scripts/integration-tests.sh
@@ -133,12 +133,20 @@ echo ""
 # Set environment variables
 export LLAMA_STACK_CLIENT_TIMEOUT=300
 
-# Setup-specific configuration is now handled by pytest via --setup
+THIS_DIR=$(dirname "$0")
+
 if [[ -n "$TEST_SETUP" ]]; then
     EXTRA_PARAMS="--setup=$TEST_SETUP"
 fi
 
-THIS_DIR=$(dirname "$0")
+# Apply setup-specific environment variables (needed for server startup and tests)
+echo "=== Applying Setup Environment Variables ==="
+SETUP_ENV=$(PYTHONPATH=$THIS_DIR/.. uv run python "$THIS_DIR/get_setup_env.py" --suite "$TEST_SUITE" --setup "$TEST_SETUP" --format bash)
+echo "Setting up environment variables:"
+echo "$SETUP_ENV"
+eval "$SETUP_ENV"
+echo ""
+
 ROOT_DIR="$THIS_DIR/.."
 cd $ROOT_DIR
 
@@ -157,6 +165,18 @@ fi
 
 # Start Llama Stack Server if needed
 if [[ "$STACK_CONFIG" == *"server:"* ]]; then
+    stop_server() {
+        echo "Stopping Llama Stack Server..."
+        pids=$(lsof -i :8321 | awk 'NR>1 {print $2}')
+        if [[ -n "$pids" ]]; then
+            echo "Killing Llama Stack Server processes: $pids"
+            kill -9 $pids
+        else
+            echo "No Llama Stack Server processes found ?!"
+        fi
+        echo "Llama Stack Server stopped"
+    }
+
     # check if server is already running
     if curl -s http://localhost:8321/v1/health 2>/dev/null | grep -q "OK"; then
         echo "Llama Stack Server is already running, skipping start"
@@ -180,6 +200,8 @@ if [[ "$STACK_CONFIG" == *"server:"* ]]; then
         done
         echo ""
     fi
+
+    trap stop_server EXIT ERR INT TERM
 fi
 
 # Run tests
@@ -239,8 +261,8 @@ pytest -s -v $PYTEST_TARGET \
     --color=yes \
     --capture=tee-sys
 exit_code=$?
-set -e
 set +x
+set -e
 
 if [ $exit_code -eq 0 ]; then
     echo "✅ All tests completed successfully"
@@ -257,18 +279,5 @@ echo "=== System Resources After Tests ==="
 free -h 2>/dev/null || echo "free command not available"
 df -h
 
-# stop server
-if [[ "$STACK_CONFIG" == *"server:"* ]]; then
-    echo "Stopping Llama Stack Server..."
-    pids=$(lsof -i :8321 | awk 'NR>1 {print $2}')
-    if [[ -n "$pids" ]]; then
-        echo "Killing Llama Stack Server processes: $pids"
-        kill -9 $pids
-    else
-        echo "No Llama Stack Server processes found ?!"
-    fi
-    echo "Llama Stack Server stopped"
-fi
-
 echo ""
 echo "=== Integration Tests Complete ==="