From 77fe91dca003df0998f0310c0b07de69c1ac64b0 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Mon, 9 Jun 2025 15:15:48 -0400 Subject: [PATCH] internal/mcp: iterators return pointers Iterators return pointers instead of values, to be consistent. Move tests out of example file. Factor out iterator tests. Change-Id: Icb14edf99d738de9d9dc08c1169057cb3dc5894b Reviewed-on: https://go-review.googlesource.com/c/tools/+/680295 Reviewed-by: Robert Findley LUCI-TryBot-Result: Go LUCI Reviewed-by: Sam Thanawalla --- mcp/client.go | 23 +++---- mcp/client_list_test.go | 136 +++++++++++++++++++++++++++++++++++++ mcp/client_test.go | 22 ++---- mcp/server_example_test.go | 112 ------------------------------ 4 files changed, 154 insertions(+), 139 deletions(-) create mode 100644 mcp/client_list_test.go diff --git a/mcp/client.go b/mcp/client.go index 15d4332..5ec7eeb 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -336,9 +336,9 @@ func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, para // Tools provides an iterator for all tools available on the server, // automatically fetching pages and managing cursors. -// The `params` argument can set the initial cursor. +// The params argument can set the initial cursor. // Iteration stops at the first encountered error, which will be yielded. -func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[Tool, error] { +func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[*Tool, error] { if params == nil { params = &ListToolsParams{} } @@ -349,9 +349,9 @@ func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) ite // Resources provides an iterator for all resources available on the server, // automatically fetching pages and managing cursors. -// The `params` argument can set the initial cursor. +// The params argument can set the initial cursor. // Iteration stops at the first encountered error, which will be yielded. -func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesParams) iter.Seq2[Resource, error] { +func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesParams) iter.Seq2[*Resource, error] { if params == nil { params = &ListResourcesParams{} } @@ -364,7 +364,7 @@ func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesPar // automatically fetching pages and managing cursors. // The `params` argument can set the initial cursor. // Iteration stops at the first encountered error, which will be yielded. -func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) iter.Seq2[ResourceTemplate, error] { +func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) iter.Seq2[*ResourceTemplate, error] { if params == nil { params = &ListResourceTemplatesParams{} } @@ -375,9 +375,9 @@ func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListReso // Prompts provides an iterator for all prompts available on the server, // automatically fetching pages and managing cursors. -// The `params` argument can set the initial cursor. +// The params argument can set the initial cursor. // Iteration stops at the first encountered error, which will be yielded. -func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) iter.Seq2[Prompt, error] { +func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) iter.Seq2[*Prompt, error] { if params == nil { params = &ListPromptsParams{} } @@ -387,17 +387,16 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) } // paginate is a generic helper function to provide a paginated iterator. -func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[T, error] { - return func(yield func(T, error) bool) { +func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[*T, error] { + return func(yield func(*T, error) bool) { for { res, err := listFunc(ctx, params) if err != nil { - var zero T - yield(zero, err) + yield(nil, err) return } for _, r := range items(res) { - if !yield(*r, nil) { + if !yield(r, nil) { return } } diff --git a/mcp/client_list_test.go b/mcp/client_list_test.go new file mode 100644 index 0000000..6d2c2c9 --- /dev/null +++ b/mcp/client_list_test.go @@ -0,0 +1,136 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp_test + +import ( + "context" + "iter" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/jsonschema" +) + +func TestListTools(t *testing.T) { + toolA := mcp.NewTool("apple", "apple tool", SayHi) + toolB := mcp.NewTool("banana", "banana tool", SayHi) + toolC := mcp.NewTool("cherry", "cherry tool", SayHi) + tools := []*mcp.ServerTool{toolA, toolB, toolC} + + wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool} + ctx := context.Background() + clientSession, serverSession, server := createSessions(ctx) + defer clientSession.Close() + defer serverSession.Close() + server.AddTools(tools...) + t.Run("ListTools", func(t *testing.T) { + res, err := clientSession.ListTools(ctx, nil) + if err != nil { + t.Fatal("ListTools() failed:", err) + } + if diff := cmp.Diff(wantTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListTools() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("ToolsIterator", func(t *testing.T) { + testIterator(ctx, t, clientSession.Tools(ctx, nil), wantTools) + }) +} + +func TestListResources(t *testing.T) { + resourceA := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://apple"}} + resourceB := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://banana"}} + resourceC := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://cherry"}} + wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource} + + resources := []*mcp.ServerResource{resourceA, resourceB, resourceC} + ctx := context.Background() + clientSession, serverSession, server := createSessions(ctx) + defer clientSession.Close() + defer serverSession.Close() + server.AddResources(resources...) + t.Run("ListResources", func(t *testing.T) { + res, err := clientSession.ListResources(ctx, nil) + if err != nil { + t.Fatal("ListResources() failed:", err) + } + if diff := cmp.Diff(wantResources, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("ResourcesIterator", func(t *testing.T) { + testIterator(ctx, t, clientSession.Resources(ctx, nil), wantResources) + }) +} + +func TestListResourceTemplates(t *testing.T) { + resourceTmplA := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://apple/{x}"}} + resourceTmplB := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://banana/{x}"}} + resourceTmplC := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://cherry/{x}"}} + wantResourceTemplates := []*mcp.ResourceTemplate{ + resourceTmplA.ResourceTemplate, resourceTmplB.ResourceTemplate, + resourceTmplC.ResourceTemplate, + } + resourceTemplates := []*mcp.ServerResourceTemplate{resourceTmplA, resourceTmplB, resourceTmplC} + ctx := context.Background() + clientSession, serverSession, server := createSessions(ctx) + defer clientSession.Close() + defer serverSession.Close() + server.AddResourceTemplates(resourceTemplates...) + t.Run("ListResourceTemplates", func(t *testing.T) { + res, err := clientSession.ListResourceTemplates(ctx, nil) + if err != nil { + t.Fatal("ListResourceTemplates() failed:", err) + } + if diff := cmp.Diff(wantResourceTemplates, res.ResourceTemplates, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListResourceTemplates() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("ResourceTemplatesIterator", func(t *testing.T) { + testIterator(ctx, t, clientSession.ResourceTemplates(ctx, nil), wantResourceTemplates) + }) +} + +func TestListPrompts(t *testing.T) { + promptA := mcp.NewPrompt("apple", "apple prompt", testPromptHandler[struct{}]) + promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}]) + promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}]) + wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt} + + prompts := []*mcp.ServerPrompt{promptA, promptB, promptC} + ctx := context.Background() + clientSession, serverSession, server := createSessions(ctx) + defer clientSession.Close() + defer serverSession.Close() + server.AddPrompts(prompts...) + t.Run("ListPrompts", func(t *testing.T) { + res, err := clientSession.ListPrompts(ctx, nil) + if err != nil { + t.Fatal("ListPrompts() failed:", err) + } + if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("PromptsIterator", func(t *testing.T) { + testIterator(ctx, t, clientSession.Prompts(ctx, nil), wantPrompts) + }) +} + +func testIterator[T any](ctx context.Context, t *testing.T, seq iter.Seq2[*T, error], want []*T) { + t.Helper() + var got []*T + for x, err := range seq { + if err != nil { + t.Fatalf("iteration failed: %v", err) + } + got = append(got, x) + } + if diff := cmp.Diff(want, got, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("mismatch (-want +got):\n%s", diff) + } +} diff --git a/mcp/client_test.go b/mcp/client_test.go index 9d24a92..757e8ee 100644 --- a/mcp/client_test.go +++ b/mcp/client_test.go @@ -50,14 +50,6 @@ var allItems = []*Item{ {"kilo", "val-K"}, } -func toItemValueSlice(ptrSlice []*Item) []Item { - var valSlice []Item - for _, ptr := range ptrSlice { - valSlice = append(valSlice, *ptr) - } - return valSlice -} - // generatePaginatedResults is a helper to create a sequence of mock responses for pagination. // It simulates a server returning items in pages based on a given page size. func generatePaginatedResults(all []*Item, pageSize int) []*ListTestResult { @@ -88,18 +80,18 @@ func TestClientPaginateBasic(t *testing.T) { results []*ListTestResult mockError error initialParams *ListTestParams - expected []Item + expected []*Item expectError bool }{ { name: "SinglePageAllItems", results: generatePaginatedResults(allItems, len(allItems)), - expected: toItemValueSlice(allItems), + expected: allItems, }, { name: "MultiplePages", results: generatePaginatedResults(allItems, 3), - expected: toItemValueSlice(allItems), + expected: allItems, }, { name: "EmptyResults", @@ -117,7 +109,7 @@ func TestClientPaginateBasic(t *testing.T) { name: "InitialCursorProvided", initialParams: &ListTestParams{Cursor: "cursor_2"}, results: generatePaginatedResults(allItems[2:], 3), - expected: toItemValueSlice(allItems[2:]), + expected: allItems[2:], }, { name: "CursorBeyondAllItems", @@ -148,7 +140,7 @@ func TestClientPaginateBasic(t *testing.T) { params = &ListTestParams{} } - var gotItems []Item + var gotItems []*Item var iterationErr error seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items }) for item, err := range seq { @@ -185,7 +177,7 @@ func TestClientPaginateVariousPageSizes(t *testing.T) { results = results[1:] return res, nil } - var gotItems []Item + var gotItems []*Item seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items }) for item, err := range seq { if err != nil { @@ -193,7 +185,7 @@ func TestClientPaginateVariousPageSizes(t *testing.T) { } gotItems = append(gotItems, item) } - if diff := cmp.Diff(toItemValueSlice(allItems), gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + if diff := cmp.Diff(allItems, gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { t.Fatalf("paginate() mismatch (-want +got):\n%s", diff) } }) diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index e233c23..3f0e40a 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -8,12 +8,8 @@ import ( "context" "fmt" "log" - "testing" - "github.com/google/go-cmp/cmp" - "github.com/google/go-cmp/cmp/cmpopts" "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/modelcontextprotocol/go-sdk/jsonschema" ) type SayHiParams struct { @@ -76,111 +72,3 @@ func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession } return clientSession, serverSession, server } - -func TestListTools(t *testing.T) { - toolA := mcp.NewTool("apple", "apple tool", SayHi) - toolB := mcp.NewTool("banana", "banana tool", SayHi) - toolC := mcp.NewTool("cherry", "cherry tool", SayHi) - tools := []*mcp.ServerTool{toolA, toolB, toolC} - ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx) - defer clientSession.Close() - defer serverSession.Close() - server.AddTools(tools...) - t.Run("ListTools", func(t *testing.T) { - wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool} - res, err := clientSession.ListTools(ctx, nil) - if err != nil { - t.Fatal("ListTools() failed:", err) - } - if diff := cmp.Diff(wantTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("ListTools() mismatch (-want +got):\n%s", diff) - } - }) - t.Run("ToolsIterator", func(t *testing.T) { - wantTools := []mcp.Tool{*toolA.Tool, *toolB.Tool, *toolC.Tool} - var gotTools []mcp.Tool - seq := clientSession.Tools(ctx, nil) - for tool, err := range seq { - if err != nil { - t.Fatalf("Tools() failed: %v", err) - } - gotTools = append(gotTools, tool) - } - if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("Tools() mismatch (-want +got):\n%s", diff) - } - }) -} - -func TestListResources(t *testing.T) { - resourceA := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://apple"}} - resourceB := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://banana"}} - resourceC := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://cherry"}} - resources := []*mcp.ServerResource{resourceA, resourceB, resourceC} - ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx) - defer clientSession.Close() - defer serverSession.Close() - server.AddResources(resources...) - t.Run("ListResources", func(t *testing.T) { - wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource} - res, err := clientSession.ListResources(ctx, nil) - if err != nil { - t.Fatal("ListResources() failed:", err) - } - if diff := cmp.Diff(wantResources, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff) - } - }) - t.Run("ResourcesIterator", func(t *testing.T) { - wantResources := []mcp.Resource{*resourceA.Resource, *resourceB.Resource, *resourceC.Resource} - var gotResources []mcp.Resource - seq := clientSession.Resources(ctx, nil) - for resource, err := range seq { - if err != nil { - t.Fatalf("Resources() failed: %v", err) - } - gotResources = append(gotResources, resource) - } - if diff := cmp.Diff(wantResources, gotResources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("Resources() mismatch (-want +got):\n%s", diff) - } - }) -} - -func TestListPrompts(t *testing.T) { - promptA := mcp.NewPrompt("apple", "apple prompt", testPromptHandler[struct{}]) - promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}]) - promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}]) - prompts := []*mcp.ServerPrompt{promptA, promptB, promptC} - ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx) - defer clientSession.Close() - defer serverSession.Close() - server.AddPrompts(prompts...) - t.Run("ListPrompts", func(t *testing.T) { - wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt} - res, err := clientSession.ListPrompts(ctx, nil) - if err != nil { - t.Fatal("ListPrompts() failed:", err) - } - if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff) - } - }) - t.Run("PromptsIterator", func(t *testing.T) { - wantPrompts := []mcp.Prompt{*promptA.Prompt, *promptB.Prompt, *promptC.Prompt} - var gotPrompts []mcp.Prompt - seq := clientSession.Prompts(ctx, nil) - for prompt, err := range seq { - if err != nil { - t.Fatalf("Prompts() failed: %v", err) - } - gotPrompts = append(gotPrompts, prompt) - } - if diff := cmp.Diff(wantPrompts, gotPrompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("Prompts() mismatch (-want +got):\n%s", diff) - } - }) -}