Skip to content
30 changes: 27 additions & 3 deletions mcptest/mcptest.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,10 @@ import (
type Server struct {
name string

tools []server.ServerTool
prompts []server.ServerPrompt
resources []server.ServerResource
tools []server.ServerTool
prompts []server.ServerPrompt
resources []server.ServerResource
resourceTemplates []ServerResourceTemplate

cancel func()

Expand Down Expand Up @@ -106,6 +107,25 @@ func (s *Server) AddResources(resources ...server.ServerResource) {
s.resources = append(s.resources, resources...)
}

// ServerResourceTemplate combines a ResourceTemplate with its handler function.
type ServerResourceTemplate struct {
Template mcp.ResourceTemplate
Handler server.ResourceTemplateHandlerFunc
}

// AddResourceTemplate adds a resource template to an unstarted server.
func (s *Server) AddResourceTemplate(template mcp.ResourceTemplate, handler server.ResourceTemplateHandlerFunc) {
s.resourceTemplates = append(s.resourceTemplates, ServerResourceTemplate{
Template: template,
Handler: handler,
})
}

// AddResourceTemplates adds multiple resource templates to an unstarted server.
func (s *Server) AddResourceTemplates(templates ...ServerResourceTemplate) {
s.resourceTemplates = append(s.resourceTemplates, templates...)
}

// Start starts the server in a goroutine. Make sure to defer Close() after Start().
// When using NewServer(), the returned server is already started.
func (s *Server) Start(ctx context.Context) error {
Expand All @@ -122,6 +142,10 @@ func (s *Server) Start(ctx context.Context) error {
mcpServer.AddTools(s.tools...)
mcpServer.AddPrompts(s.prompts...)
mcpServer.AddResources(s.resources...)

for _, template := range s.resourceTemplates {
mcpServer.AddResourceTemplate(template.Template, template.Handler)
}

logger := log.New(&s.logBuffer, "", 0)

Expand Down
76 changes: 76 additions & 0 deletions mcptest/mcptest_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -187,3 +187,79 @@ func TestServerWithResource(t *testing.T) {
t.Errorf("Got %q, want %q", textContent.Text, want)
}
}

func TestServerWithResourceTemplate(t *testing.T) {
ctx := context.Background()

srv := mcptest.NewUnstartedServer(t)
defer srv.Close()

template := mcp.NewResourceTemplate(
"file://users/{userId}/documents/{docId}",
"User Document",
mcp.WithTemplateDescription("A user's document"),
mcp.WithTemplateMIMEType("text/plain"),
)

handler := func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
if request.Params.Arguments == nil {
return nil, fmt.Errorf("expected arguments to be populated from URI template")
}

userIds, ok := request.Params.Arguments["userId"].([]string)
if !ok {
return nil, fmt.Errorf("expected userId argument to be populated from URI template")
}
if len(userIds) != 1 {
return nil, fmt.Errorf("expected userId to have one value, but got %d", len(userIds))
}
if userIds[0] != "john" {
return nil, fmt.Errorf("expected userId argument to be 'john', got %s", userIds[0])
}

docIds, ok := request.Params.Arguments["docId"].([]string)
if !ok {
return nil, fmt.Errorf("expected docId argument to be populated from URI template")
}
if len(docIds) != 1 {
return nil, fmt.Errorf("expected docId to have one value, but got %d", len(docIds))
}
if docIds[0] != "readme.txt" {
return nil, fmt.Errorf("expected docId argument to be 'readme.txt', got %v", docIds)
}

return []mcp.ResourceContents{
mcp.TextResourceContents{
URI: request.Params.URI,
MIMEType: "text/plain",
Text: fmt.Sprintf("Document %s for user %s", docIds[0], userIds[0]),
},
}, nil
}

srv.AddResourceTemplate(template, handler)

err := srv.Start(ctx)
if err != nil {
t.Fatal(err)
}

// Test reading a resource that matches the template
var readReq mcp.ReadResourceRequest
readReq.Params.URI = "file://users/john/documents/readme.txt"
readResult, err := srv.Client().ReadResource(ctx, readReq)
if err != nil {
t.Fatal("ReadResource:", err)
}
if len(readResult.Contents) != 1 {
t.Fatalf("Expected 1 content, got %d", len(readResult.Contents))
}
textContent, ok := readResult.Contents[0].(mcp.TextResourceContents)
if !ok {
t.Fatalf("Expected TextResourceContents, got %T", readResult.Contents[0])
}
want := "Document readme.txt for user john"
if textContent.Text != want {
t.Errorf("Got %q, want %q", textContent.Text, want)
}
}