From 735831206d66185fbe8653291abd9132f1a292ac Mon Sep 17 00:00:00 2001
From: Omar Abdelwahab <omara@fb.com>
Date: Fri, 7 Nov 2025 14:46:30 -0800
Subject: [PATCH] fix: update tests to use new mcp_authorization field

Updates integration tests to use the new mcp_authorization field
instead of the old method of passing Authorization in mcp_headers.

Changes:
- tests/integration/tool_runtime/test_mcp.py
- tests/integration/inference/test_tools_with_schemas.py
- tests/integration/tool_runtime/test_mcp_json_schema.py (6 occurrences)

All tests now use:
  provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}

Instead of the old rejected format:
  provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}

This aligns with the security architecture that prevents
accidentally leaking inference tokens to MCP servers.
---
 .../integration/inference/test_tools_with_schemas.py |  2 +-
 tests/integration/tool_runtime/test_mcp.py           |  6 ++----
 .../integration/tool_runtime/test_mcp_json_schema.py | 12 ++++++------
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/tests/integration/inference/test_tools_with_schemas.py b/tests/integration/inference/test_tools_with_schemas.py
index f30e9ece5..9a3ac0bf0 100644
--- a/tests/integration/inference/test_tools_with_schemas.py
+++ b/tests/integration/inference/test_tools_with_schemas.py
@@ -193,7 +193,7 @@ class TestMCPToolsInChatCompletion:
             mcp_endpoint=dict(uri=uri),
         )
 
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }
diff --git a/tests/integration/tool_runtime/test_mcp.py b/tests/integration/tool_runtime/test_mcp.py
index a3939fd0a..8138f0d92 100644
--- a/tests/integration/tool_runtime/test_mcp.py
+++ b/tests/integration/tool_runtime/test_mcp.py
@@ -43,10 +43,8 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
     )
 
     provider_data = {
-        "mcp_headers": {
-            uri: {
-                "Authorization": f"Bearer {AUTH_TOKEN}",
-            },
+        "mcp_authorization": {
+            uri: AUTH_TOKEN,  # Token
         },
     }
     auth_headers = {
diff --git a/tests/integration/tool_runtime/test_mcp_json_schema.py b/tests/integration/tool_runtime/test_mcp_json_schema.py
index def0b27b8..6302fa385 100644
--- a/tests/integration/tool_runtime/test_mcp_json_schema.py
+++ b/tests/integration/tool_runtime/test_mcp_json_schema.py
@@ -123,7 +123,7 @@ class TestMCPSchemaPreservation:
             mcp_endpoint=dict(uri=uri),
         )
 
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }
@@ -166,7 +166,7 @@ class TestMCPSchemaPreservation:
             provider_id="model-context-protocol",
             mcp_endpoint=dict(uri=uri),
         )
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }
@@ -216,7 +216,7 @@ class TestMCPSchemaPreservation:
             mcp_endpoint=dict(uri=uri),
         )
 
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }
@@ -263,7 +263,7 @@ class TestMCPToolInvocation:
             mcp_endpoint=dict(uri=uri),
         )
 
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }
@@ -309,7 +309,7 @@ class TestMCPToolInvocation:
             mcp_endpoint=dict(uri=uri),
         )
 
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }
@@ -365,7 +365,7 @@ class TestAgentWithMCPTools:
             mcp_endpoint=dict(uri=uri),
         )
 
-        provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
+        provider_data = {"mcp_authorization": {uri: AUTH_TOKEN}}  # Token without "Bearer " prefix
         auth_headers = {
             "X-LlamaStack-Provider-Data": json.dumps(provider_data),
         }