diff --git a/tests/integration/inference/test_tools_with_schemas.py b/tests/integration/inference/test_tools_with_schemas.py
index 5b6e69ae3..8ceea3930 100644
--- a/tests/integration/inference/test_tools_with_schemas.py
+++ b/tests/integration/inference/test_tools_with_schemas.py
@@ -183,11 +183,11 @@ class TestMCPToolsInChatCompletion:
         uri = mcp_with_schemas["server_url"]
 
         try:
-            llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
+            llama_stack_client.toolgroups.unregister_toolgroup(toolgroup_id=test_toolgroup_id)
         except Exception:
             pass
 
-        llama_stack_client.toolgroups.register(
+        llama_stack_client.toolgroups.register_tool_group(
             toolgroup_id=test_toolgroup_id,
             provider_id="model-context-protocol",
             mcp_endpoint=dict(uri=uri),