phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-05 10:13:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

update test_register_and_unregister_toolgroup

Browse source 
Signed-off-by: Paolo Dettori <dettori@us.ibm.com>
This commit is contained in:
Paolo Dettori 2025-04-07 21:48:01 -04:00
parent da44d78b99
commit 0c3f9f46f5
1 changed files with 26 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

42
tests/integration/tool_runtime/test_registration.py
View file

@ -4,6 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import socket
import threading import threading
import time import time
@ -48,8 +49,15 @@ def mcp_server():
], ],
) )
def get_open_port():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("", 0))
return sock.getsockname()[1]
port = get_open_port()
def run_server(): def run_server():
uvicorn.run(app, host="0.0.0.0", port=8000) uvicorn.run(app, host="0.0.0.0", port=port)
# Start the server in a new thread # Start the server in a new thread
server_thread = threading.Thread(target=run_server, daemon=True) server_thread = threading.Thread(target=run_server, daemon=True)
@ -61,54 +69,56 @@ def mcp_server():
while time.time() - start_time < timeout: while time.time() - start_time < timeout:
try: try:
response = httpx.get("http://localhost:8000/sse") response = httpx.get(f"http://localhost:{port}/sse")
if response.status_code == 200: if response.status_code == 200:
break break
except (httpx.RequestError, httpx.HTTPStatusError): except (httpx.RequestError, httpx.HTTPStatusError):
pass pass
time.sleep(0.1) time.sleep(0.1)
yield yield port
def test_register_and_unregister_toolgroup(client_with_models, mcp_server): def test_register_and_unregister_toolgroup(llama_stack_client, mcp_server):
""" """
Integration test for registering and unregistering a toolgroup using the ToolGroups API. Integration test for registering and unregistering a toolgroup using the ToolGroups API.
""" """
port = mcp_server
test_toolgroup_id = "remote::web-fetch" test_toolgroup_id = "remote::web-fetch"
provider_id = "model-context-protocol" provider_id = "model-context-protocol"
# Cleanup before running the test # Cleanup before running the test
test_toolgroup = client_with_models.toolgroups.get(toolgroup_id=test_toolgroup_id) toolgroups = llama_stack_client.toolgroups.list()
if test_toolgroup is not None: for toolgroup in toolgroups:
client_with_models.toolgroups.unregister(toolgroup_id=test_toolgroup_id) if toolgroup.identifier == test_toolgroup_id:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
# Register the toolgroup # Register the toolgroup
client_with_models.toolgroups.register( llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id, toolgroup_id=test_toolgroup_id,
provider_id=provider_id, provider_id=provider_id,
mcp_endpoint=URL(uri="http://localhost:8000/sse"), mcp_endpoint=URL(uri=f"http://localhost:{port}/sse"),
) )
# Verify registration # Verify registration
registered_toolgroup = client_with_models.toolgroups.get(toolgroup_id=test_toolgroup_id) registered_toolgroup = llama_stack_client.toolgroups.get(toolgroup_id=test_toolgroup_id)
assert registered_toolgroup is not None assert registered_toolgroup is not None
assert registered_toolgroup.identifier == test_toolgroup_id assert registered_toolgroup.identifier == test_toolgroup_id
assert registered_toolgroup.provider_id == provider_id assert registered_toolgroup.provider_id == provider_id
# Verify tools listing # Verify tools listing
tools_list_response = client_with_models.tools.list(toolgroup_id=test_toolgroup_id) tools_list_response = llama_stack_client.tools.list(toolgroup_id=test_toolgroup_id)
assert isinstance(tools_list_response, list) assert isinstance(tools_list_response, list)
assert tools_list_response assert tools_list_response
# Unregister the toolgroup # Unregister the toolgroup
client_with_models.toolgroups.unregister(toolgroup_id=test_toolgroup_id) llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
# Verify unregistration # Verify it is unregistered
unregistered_toolgroup = client_with_models.toolgroups.get(toolgroup_id=test_toolgroup_id) with pytest.raises(ValueError, match=f"Tool group '{test_toolgroup_id}' not found"):
assert unregistered_toolgroup is None llama_stack_client.toolgroups.get(toolgroup_id=test_toolgroup_id)
# Verify tools are also unregistered # Verify tools are also unregistered
unregister_tools_list_response = client_with_models.tools.list(toolgroup_id=test_toolgroup_id) unregister_tools_list_response = llama_stack_client.tools.list(toolgroup_id=test_toolgroup_id)
assert isinstance(unregister_tools_list_response, list) assert isinstance(unregister_tools_list_response, list)
assert not unregister_tools_list_response assert not unregister_tools_list_response