mcp.go•1.87 kB
package test
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/mark3labs/mcp-go/client"
"github.com/mark3labs/mcp-go/client/transport"
"github.com/mark3labs/mcp-go/mcp"
"github.com/stretchr/testify/require"
)
func McpInitRequest() mcp.InitializeRequest {
initRequest := mcp.InitializeRequest{
Request: mcp.Request{Method: "initialize"},
}
initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
initRequest.Params.ClientInfo = mcp.Implementation{Name: "test", Version: "1.33.7"}
return initRequest
}
type McpClient struct {
ctx context.Context
testServer *httptest.Server
*client.Client
}
func NewMcpClient(t *testing.T, mcpHttpServer http.Handler, options ...transport.StreamableHTTPCOption) *McpClient {
require.NotNil(t, mcpHttpServer, "McpHttpServer must be provided")
var err error
ret := &McpClient{ctx: t.Context()}
ret.testServer = httptest.NewServer(mcpHttpServer)
options = append(options, transport.WithContinuousListening())
ret.Client, err = client.NewStreamableHttpClient(ret.testServer.URL+"/mcp", options...)
require.NoError(t, err, "Expected no error creating MCP client")
err = ret.Start(t.Context())
require.NoError(t, err, "Expected no error starting MCP client")
_, err = ret.Initialize(t.Context(), McpInitRequest())
require.NoError(t, err, "Expected no error initializing MCP client")
return ret
}
func (m *McpClient) Close() {
if m.Client != nil {
_ = m.Client.Close()
}
if m.testServer != nil {
m.testServer.Close()
}
}
// CallTool helper function to call a tool by name with arguments
func (m *McpClient) CallTool(name string, args map[string]interface{}) (*mcp.CallToolResult, error) {
callToolRequest := mcp.CallToolRequest{}
callToolRequest.Params.Name = name
callToolRequest.Params.Arguments = args
return m.Client.CallTool(m.ctx, callToolRequest)
}