Set test client info (#692)

* allow test to set client info

This allows testing of mcp server behavior that depends on the type of
the mcp client.

* test allowing test to set client info

* CR comment check Start() return code and use ctx consistently

---------

Co-authored-by: Uri Simchoni <uri@vfunction.com>

Author: urisimchoni

mcptest/mcptest.go

@@ -24,6 +24,7 @@ type Server struct {
 	prompts           []server.ServerPrompt
 	resources         []server.ServerResource
 	resourceTemplates []server.ServerResourceTemplate
+	clientInfo        mcp.Implementation
 
 	cancel func()
 
@@ -120,6 +121,11 @@ func (s *Server) AddResourceTemplates(templates ...server.ServerResourceTemplate
 	s.resourceTemplates = append(s.resourceTemplates, templates...)
 }
 
+// SetClientInfo sets the client info for the test client.
+func (s *Server) SetClientInfo(info mcp.Implementation) {
+	s.clientInfo = info
+}
+
 // 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 {
@@ -157,6 +163,7 @@ func (s *Server) Start(ctx context.Context) error {
 
 	var initReq mcp.InitializeRequest
 	initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
+	initReq.Params.ClientInfo = s.clientInfo
 	if _, err := s.client.Initialize(ctx, initReq); err != nil {
 		return fmt.Errorf("client.Initialize(): %w", err)
 	}

mcptest/mcptest_test.go

@@ -411,3 +411,47 @@ func TestListToolsWithHeader(t *testing.T) {
 		t.Fatalf("Expected value is %s, got %s", expectedHeaderValue, gotHeaderValue)
 	}
 }
+
+func TestSimulateClientInfo(t *testing.T) {
+	ctx := context.Background()
+
+	srv := mcptest.NewUnstartedServer(t)
+	defer srv.Close()
+	srv.AddTool(mcp.NewTool("whoami", mcp.WithDescription("Says hello to client.")),
+		func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			clientName := ""
+			if clientSession := server.ClientSessionFromContext(ctx); clientSession != nil {
+				if sessionWithClientInfo, ok := clientSession.(server.SessionWithClientInfo); ok {
+					clientName = sessionWithClientInfo.GetClientInfo().Name
+				}
+			}
+			return mcp.NewToolResultText(fmt.Sprintf("Hello, %s!", clientName)), nil
+		})
+	srv.SetClientInfo(mcp.Implementation{
+		Name: "test-client",
+	})
+	err := srv.Start(ctx)
+	if err != nil {
+		t.Fatal("Start:", err)
+	}
+
+	client := srv.Client()
+
+	var req mcp.CallToolRequest
+	req.Params.Name = "whoami"
+
+	result, err := client.CallTool(ctx, req)
+	if err != nil {
+		t.Fatal("CallTool:", err)
+	}
+
+	got, err := resultToString(result)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want := "Hello, test-client!"
+	if got != want {
+		t.Errorf("Got %q, want %q", got, want)
+	}
+}