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 <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
This commit is contained in:
Jonathan Amsterdam 2025-06-09 15:15:48 -04:00
commit 77fe91dca0

View file

@ -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)
}
})
}