|
| 1 | +// Package mcptest implements helper functions for testing MCP servers. |
| 2 | +package mcptest |
| 3 | + |
| 4 | +import ( |
| 5 | + "bytes" |
| 6 | + "context" |
| 7 | + "io" |
| 8 | + "log" |
| 9 | + "testing" |
| 10 | + |
| 11 | + "github.com/mark3labs/mcp-go/client" |
| 12 | + "github.com/mark3labs/mcp-go/mcp" |
| 13 | + "github.com/mark3labs/mcp-go/server" |
| 14 | +) |
| 15 | + |
| 16 | +// Server encapsulates an MCP server and manages resources like pipes and context. |
| 17 | +type Server struct { |
| 18 | + name string |
| 19 | + tools []server.ServerTool |
| 20 | + |
| 21 | + ctx context.Context |
| 22 | + cancel func() |
| 23 | + |
| 24 | + serverReader io.Reader |
| 25 | + serverWriter io.Writer |
| 26 | + clientReader io.Reader |
| 27 | + clientWriter io.WriteCloser |
| 28 | + |
| 29 | + logBuffer bytes.Buffer |
| 30 | + |
| 31 | + client *client.StdioMCPClient |
| 32 | +} |
| 33 | + |
| 34 | +// NewServer starts a new MCP server with the provided tools and returns the server instance. |
| 35 | +func NewServer(t *testing.T, tools ...server.ServerTool) *Server { |
| 36 | + server := NewUnstartedServer(t) |
| 37 | + server.AddTools(tools...) |
| 38 | + server.Start() |
| 39 | + |
| 40 | + return server |
| 41 | +} |
| 42 | + |
| 43 | +// NewUnstartedServer creates a new MCP server instance with the given name, but does not start the server. |
| 44 | +// Useful for tests where you need to add tools before starting the server. |
| 45 | +func NewUnstartedServer(t *testing.T) *Server { |
| 46 | + server := &Server{ |
| 47 | + name: t.Name(), |
| 48 | + } |
| 49 | + |
| 50 | + // TODO: use t.Context() once we switch to go >= 1.24 |
| 51 | + ctx := context.Background() |
| 52 | + |
| 53 | + // Set up context with cancellation, used to stop the server |
| 54 | + server.ctx, server.cancel = context.WithCancel(ctx) |
| 55 | + |
| 56 | + // Set up pipes for client-server communication |
| 57 | + server.serverReader, server.clientWriter = io.Pipe() |
| 58 | + server.clientReader, server.serverWriter = io.Pipe() |
| 59 | + |
| 60 | + // Return the configured server |
| 61 | + return server |
| 62 | +} |
| 63 | + |
| 64 | +// AddTools adds multiple tools to an unstarted server. |
| 65 | +func (s *Server) AddTools(tools ...server.ServerTool) { |
| 66 | + s.tools = append(s.tools, tools...) |
| 67 | +} |
| 68 | + |
| 69 | +// AddTool adds a tool to an unstarted server. |
| 70 | +func (s *Server) AddTool(tool mcp.Tool, handler server.ToolHandlerFunc) { |
| 71 | + s.tools = append(s.tools, server.ServerTool{ |
| 72 | + Tool: tool, |
| 73 | + Handler: handler, |
| 74 | + }) |
| 75 | +} |
| 76 | + |
| 77 | +// Start starts the server in a goroutine. Make sure to defer Close() after Start(). |
| 78 | +// When using NewServer(), the returned server is already started. |
| 79 | +func (s *Server) Start() { |
| 80 | + // Start the MCP server in a goroutine |
| 81 | + go func() { |
| 82 | + mcpServer := server.NewMCPServer(s.name, "1.0.0") |
| 83 | + |
| 84 | + mcpServer.AddTools(s.tools...) |
| 85 | + |
| 86 | + logger := log.New(&s.logBuffer, "", 0) |
| 87 | + |
| 88 | + stdioServer := server.NewStdioServer(mcpServer) |
| 89 | + stdioServer.SetErrorLogger(logger) |
| 90 | + |
| 91 | + if err := stdioServer.Listen(s.ctx, s.serverReader, s.serverWriter); err != nil { |
| 92 | + logger.Println("StdioServer.Listen failed:", err) |
| 93 | + } |
| 94 | + }() |
| 95 | + |
| 96 | + s.client = client.NewStdioMCPClientWithIO(s.clientReader, s.clientWriter, io.NopCloser(&s.logBuffer)) |
| 97 | +} |
| 98 | + |
| 99 | +// Close stops the server and cleans up resources like temporary directories. |
| 100 | +func (s *Server) Close() { |
| 101 | + if s.client != nil { |
| 102 | + s.client.Close() |
| 103 | + s.client = nil |
| 104 | + } |
| 105 | + |
| 106 | + if s.cancel != nil { |
| 107 | + s.cancel() |
| 108 | + s.cancel = nil |
| 109 | + } |
| 110 | +} |
| 111 | + |
| 112 | +// Client returns an MCP client connected to the server. |
| 113 | +func (s *Server) Client() client.MCPClient { |
| 114 | + return s.client |
| 115 | +} |
0 commit comments