diff --git a/tests/unit/providers/utils/tools/test_mcp.py b/tests/unit/providers/utils/tools/test_mcp.py new file mode 100644 index 000000000..9ee90cd3a --- /dev/null +++ b/tests/unit/providers/utils/tools/test_mcp.py @@ -0,0 +1,90 @@ +import asyncio +import re +from typing import List +from mcp import ListToolsResult +import pytest +from unittest.mock import AsyncMock, patch +import json + +from llama_stack.apis.tools.tools import ToolParameter +from llama_stack.providers.utils.tools.mcp import list_mcp_tools + + +def find_param(params:List[ToolParameter], param_name: str)->ToolParameter| None: + return next((p for p in params if p.name == param_name), None) + +@pytest.mark.asyncio +async def test_list_mcp_tools_with_ref_defs(): + mcp_tools_resp_json = """ + {"tools":[{"name":"book_reservation", + "inputSchema":{ + "$defs":{ + "FlightInfo":{"properties":{"flight_number":{"description":"Flight number, such as \'HAT001\'.","title":"Flight Number","type":"string"},"date":{"description":"The date for the flight in the format \'YYYY-MM-DD\', such as \'2024-05-01\'.","title":"Date","type":"string"}},"required":["flight_number","date"],"title":"FlightInfo","type":"object"}, + "Passenger":{"properties":{"first_name":{"description":"Passenger\'s first name","title":"First Name","type":"string"},"last_name":{"description":"Passenger\'s last name","title":"Last Name","type":"string"},"dob":{"description":"Date of birth in YYYY-MM-DD format","title":"Dob","type":"string"}},"required":["first_name","last_name","dob"],"title":"Passenger","type":"object"}, + "Payment":{"properties":{"payment_id":{"description":"Unique reference for the payment method in the user\'s payment methods.","title":"Payment Id","type":"string"},"amount":{"description":"Payment amount in dollars","title":"Amount","type":"integer"}},"required":["payment_id","amount"],"title":"Payment","type":"object"} + }, + "properties":{ + "user_id":{"title":"User Id","type":"string"}, + "origin":{"title":"Origin","type":"string"}, + "destination":{"title":"Destination","type":"string"}, + "flight_type":{"enum":["round_trip","one_way"],"title":"Flight Type","type":"string"}, + "cabin":{"enum":["business","economy","basic_economy"],"title":"Cabin","type":"string"}, + "flights":{"items":{"$ref":"#/$defs/FlightInfo"},"title":"Flights","type":"array"}, + "passengers":{"items":{"anyOf":[{"$ref":"#/$defs/Passenger"},{"additionalProperties":true,"type":"object"}]},"title":"Passengers","type":"array"}, + "payment":{"$ref":"#/$defs/Payment","title":"Payment Methods"}, + "total_baggages":{"title":"Total Baggages","type":"integer"}, + "nonfree_baggages":{"title":"Nonfree Baggages","type":"integer"}, + "insurance":{"enum":["yes","no"],"title":"Insurance","type":"string"} + }, + "required":["user_id","origin","destination","flight_type","cabin","flights","passengers","payment_methods","total_baggages","nonfree_baggages","insurance"], + "type":"object" + }, + + "outputSchema":{ + "$defs":{ + "Passenger":{"properties":{"first_name":{"description":"Passenger\'s first name","title":"First Name","type":"string"},"last_name":{"description":"Passenger\'s last name","title":"Last Name","type":"string"},"dob":{"description":"Date of birth in YYYY-MM-DD format","title":"Dob","type":"string"}},"required":["first_name","last_name","dob"],"title":"Passenger","type":"object"}, + "Payment":{"properties":{"payment_id":{"description":"Unique reference for the payment method in the user\'s payment methods.","title":"Payment Id","type":"string"},"amount":{"description":"Payment amount in dollars","title":"Amount","type":"integer"}},"required":["payment_id","amount"],"title":"Payment","type":"object"}, + "ReservationFlight":{"properties":{"flight_number":{"description":"Unique flight identifier","title":"Flight Number","type":"string"},"origin":{"description":"IATA code for origin airport","title":"Origin","type":"string"},"destination":{"description":"IATA code for destination airport","title":"Destination","type":"string"},"date":{"description":"Flight date in YYYY-MM-DD format","title":"Date","type":"string"},"price":{"description":"Flight price in dollars.","title":"Price","type":"integer"}},"required":["flight_number","origin","destination","date","price"],"title":"ReservationFlight","type":"object"} + }, + "properties":{ + "reservation_id":{"description":"Unique identifier for the reservation","title":"Reservation Id","type":"string"}, + "user_id":{"description":"ID of the user who made the reservation","title":"User Id","type":"string"}, + "origin":{"description":"IATA code for trip origin","title":"Origin","type":"string"},"destination":{"description":"IATA code for trip destination","title":"Destination","type":"string"}, + "flight_type":{"description":"Type of trip","enum":["round_trip","one_way"],"title":"Flight Type","type":"string"}, + "cabin":{"description":"Selected cabin class","enum":["business","economy","basic_economy"],"title":"Cabin","type":"string"}, + "flights":{"description":"List of flights in the reservation","items":{"$ref":"#/$defs/ReservationFlight"},"title":"Flights","type":"array"},"passengers":{"description":"List of passengers on the reservation","items":{"$ref":"#/$defs/Passenger"},"title":"Passengers","type":"array"}, + "payment_history":{"description":"History of payments for this reservation","items":{"$ref":"#/$defs/Payment"},"title":"Payment History","type":"array"}, + "created_at":{"description":"Timestamp when reservation was created in the format YYYY-MM-DDTHH:MM:SS","title":"Created At","type":"string"}, + "total_baggages":{"description":"Total number of bags in reservation","title":"Total Baggages","type":"integer"}, + "nonfree_baggages":{"description":"Number of paid bags in reservation","title":"Nonfree Baggages","type":"integer"}, + "insurance":{"description":"Whether travel insurance was purchased","enum":["yes","no"],"title":"Insurance","type":"string"}, + "status":{"anyOf":[{"const":"cancelled","type":"string"},{"type":"null"}],"default":null,"description":"Status of the reservation","title":"Status"} + }, + "required":["reservation_id","user_id","origin","destination","flight_type","cabin","flights","passengers","payment_history","created_at","total_baggages","nonfree_baggages","insurance"],"title":"Reservation","type":"object"}, + "annotations":null,"meta":{"_fastmcp":{"tags":[]}} + }]}""" + mcp_tools_resp = ListToolsResult.model_validate_json(mcp_tools_resp_json) + + # Mock the client_wrapper context manager + mock_session = AsyncMock() + mock_session.list_tools.return_value = mcp_tools_resp + + mock_cm = AsyncMock() + mock_cm.__aenter__.return_value = mock_session + mock_cm.__aexit__.return_value = None + + # Patch the client_wrapper to return our mock context manager + with patch("llama_stack.providers.utils.tools.mcp.client_wrapper", return_value=mock_cm): + tools_data = await list_mcp_tools("fake_endpoint", {"Authorization": "Bearer X"}) + tools = tools_data.data + assert len(tools) == 1 + book_resv = tools[0] + assert book_resv.name == "book_reservation" + assert find_param(book_resv.parameters, "payment").properties.payment_id.type == "string" + assert find_param(book_resv.parameters, "flights").items.properties.flight_number.type == "string" + assert find_param(book_resv.parameters, "passengers").items.properties.first_name.title == "First Name" + mock_session.list_tools.assert_awaited_once() + +asyncio.run( + test_list_mcp_tools_with_ref_defs() +)