test(mcp): store initialization results in McpClient for better instruction handling (#582)

Signed-off-by: Marc Nuri <[email protected]>

Author: manusa

internal/test/mcp.go

@@ -25,6 +25,7 @@ type McpClient struct {
 	ctx        context.Context
 	testServer *httptest.Server
 	*client.Client
+	InitializeResult *mcp.InitializeResult
 }
 
 func NewMcpClient(t *testing.T, mcpHttpServer http.Handler, options ...transport.StreamableHTTPCOption) *McpClient {
@@ -37,7 +38,7 @@ func NewMcpClient(t *testing.T, mcpHttpServer http.Handler, options ...transport
 	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())
+	ret.InitializeResult, err = ret.Initialize(t.Context(), McpInitRequest())
 	require.NoError(t, err, "Expected no error initializing MCP client")
 	return ret
 }

pkg/mcp/mcp_test.go

@@ -1,14 +1,12 @@
 package mcp
 
 import (
-	"context"
 	"net/http"
-	"net/http/httptest"
 	"sync"
 	"testing"
 
+	"github.com/BurntSushi/toml"
 	"github.com/containers/kubernetes-mcp-server/internal/test"
-	"github.com/mark3labs/mcp-go/client"
 	"github.com/mark3labs/mcp-go/client/transport"
 	"github.com/stretchr/testify/suite"
 )
@@ -103,66 +101,25 @@ func TestMcpHeaders(t *testing.T) {
 
 type ServerInstructionsSuite struct {
 	BaseMcpSuite
-	mockServer *test.MockServer
 }
 
-func (s *ServerInstructionsSuite) SetupTest() {
-	s.BaseMcpSuite.SetupTest()
-	s.mockServer = test.NewMockServer()
-	s.mockServer.Handle(&test.DiscoveryClientHandler{})
-	s.Cfg.KubeConfig = s.mockServer.KubeconfigFile(s.T())
-}
-
-func (s *ServerInstructionsSuite) TearDownTest() {
-	s.BaseMcpSuite.TearDownTest()
-	if s.mockServer != nil {
-		s.mockServer.Close()
-	}
-}
-
-func (s *ServerInstructionsSuite) TestServerInstructions() {
+func (s *ServerInstructionsSuite) TestServerInstructionsEmpty() {
+	s.InitMcpClient()
 	s.Run("returns empty instructions when not configured", func() {
-		s.Cfg.ServerInstructions = ""
-		mcpServer, err := NewServer(Configuration{StaticConfig: s.Cfg})
-		s.Require().NoError(err)
-		s.T().Cleanup(mcpServer.Close)
-
-		testServer := httptest.NewServer(mcpServer.ServeHTTP())
-		s.T().Cleanup(testServer.Close)
-
-		mcpClient, err := client.NewStreamableHttpClient(testServer.URL+"/mcp", transport.WithContinuousListening())
-		s.Require().NoError(err)
-		s.T().Cleanup(func() { _ = mcpClient.Close() })
-
-		err = mcpClient.Start(context.Background())
-		s.Require().NoError(err)
-
-		result, err := mcpClient.Initialize(context.Background(), test.McpInitRequest())
-		s.Require().NoError(err)
-		s.Empty(result.Instructions, "instructions should be empty when not configured")
+		s.Require().NotNil(s.InitializeResult)
+		s.Empty(s.InitializeResult.Instructions, "instructions should be empty when not configured")
 	})
+}
 
+func (s *ServerInstructionsSuite) TestServerInstructionsFromConfiguration() {
+	s.Require().NoError(toml.Unmarshal([]byte(`
+		server_instructions = "Always use YAML output format for kubectl commands."
+	`), s.Cfg), "Expected to parse server instructions config")
+	s.InitMcpClient()
 	s.Run("returns configured instructions", func() {
-		expectedInstructions := "Always use YAML output format for kubectl commands."
-		s.Cfg.ServerInstructions = expectedInstructions
-
-		mcpServer, err := NewServer(Configuration{StaticConfig: s.Cfg})
-		s.Require().NoError(err)
-		s.T().Cleanup(mcpServer.Close)
-
-		testServer := httptest.NewServer(mcpServer.ServeHTTP())
-		s.T().Cleanup(testServer.Close)
-
-		mcpClient, err := client.NewStreamableHttpClient(testServer.URL+"/mcp", transport.WithContinuousListening())
-		s.Require().NoError(err)
-		s.T().Cleanup(func() { _ = mcpClient.Close() })
-
-		err = mcpClient.Start(context.Background())
-		s.Require().NoError(err)
-
-		result, err := mcpClient.Initialize(context.Background(), test.McpInitRequest())
-		s.Require().NoError(err)
-		s.Equal(expectedInstructions, result.Instructions, "instructions should match configured value")
+		s.Require().NotNil(s.InitializeResult)
+		s.Equal("Always use YAML output format for kubectl commands.", s.InitializeResult.Instructions,
+			"instructions should match configured value")
 	})
 }