simplification

scripts/get_setup_env.py

@@ -47,24 +47,75 @@ def get_setup_env_vars(setup_name, suite_name=None):
     return setup.env
 
 
+def get_setup_config(setup_name, suite_name=None):
+    """
+    Get full configuration (env vars + defaults) for a setup.
+
+    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 with 'env' and 'defaults' keys
+    """
+    # 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 {"env": {}, "defaults": {}}
+
+    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 {"env": setup.env, "defaults": setup.defaults}
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Extract environment variables from a test setup")
+    parser = argparse.ArgumentParser(description="Extract environment variables and defaults 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)")
+    parser.add_argument(
+        "--include-defaults",
+        action="store_true",
+        help="Include setup defaults (text_model, embedding_model, etc.) in addition to env vars",
+    )
 
     args = parser.parse_args()
 
-    env_vars = get_setup_env_vars(args.setup, args.suite)
+    if args.include_defaults:
+        config = get_setup_config(args.setup, args.suite)
+        if args.format == "bash":
+            # Output env vars as bash export statements
+            for key, value in config["env"].items():
+                print(f"export {key}='{value}'")
+            # Output defaults as bash export statements with LLAMA_STACK_TEST_ prefix
+            for key, value in config["defaults"].items():
+                env_key = f"LLAMA_STACK_TEST_{key.upper()}"
+                print(f"export {env_key}='{value}'")
+        elif args.format == "json":
+            import json
 
-    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(config))
+    else:
+        env_vars = get_setup_env_vars(args.setup, args.suite)
 
-        print(json.dumps(env_vars))
+        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__":

scripts/integration-tests.sh

@@ -198,36 +198,12 @@ echo "$SETUP_ENV"
 eval "$SETUP_ENV"
 echo ""
 if [[ -n "$RESOLVED_TEST_SETUP" ]]; then
-    SETUP_DEFAULTS=$(PYTHONPATH=$THIS_DIR/.. python - "$RESOLVED_TEST_SETUP" <<'PY'
-import sys
-from tests.integration.suites import SETUP_DEFINITIONS
+    # Export setup name - TypeScript tests will call get_setup_env.py themselves to get model defaults
+    export LLAMA_STACK_TEST_SETUP="$RESOLVED_TEST_SETUP"
 
-setup_name = sys.argv[1]
-if not setup_name:
-    sys.exit(0)
-
-setup = SETUP_DEFINITIONS.get(setup_name)
-if not setup:
-    sys.exit(0)
-
-for key, value in setup.defaults.items():
-    print(f"{key}={value}")
-PY
-)
-
-    while IFS='=' read -r key value; do
-        case "$key" in
-        text_model)
-            export LLAMA_STACK_TEST_MODEL="$value"
-            ;;
-        embedding_model)
-            export LLAMA_STACK_TEST_EMBEDDING_MODEL="$value"
-            ;;
-        vision_model)
-            export LLAMA_STACK_TEST_VISION_MODEL="$value"
-            ;;
-        esac
-    done <<< "$SETUP_DEFAULTS"
+    # Export model env vars for Python tests using get_setup_env.py
+    SETUP_DEFAULTS_ENV=$(PYTHONPATH=$THIS_DIR/.. python $THIS_DIR/get_setup_env.py --setup "$RESOLVED_TEST_SETUP" --format bash --include-defaults)
+    eval "$SETUP_DEFAULTS_ENV"
 fi
 
 ROOT_DIR="$THIS_DIR/.."