Skip to content

Commit 5da1a92

Browse files
manusamanusa
authored
test(mcp): update MCP headers tests to use testify and improve readability (containers#417)
Signed-off-by: Marc Nuri <[email protected]>
1 parent 44053f1 commit 5da1a92

File tree

3 files changed

+55
-46
lines changed

3 files changed

+55
-46
lines changed

‎internal/test/mcp.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"testing"
77

88
"github.com/mark3labs/mcp-go/client"
9+
"github.com/mark3labs/mcp-go/client/transport"
910
"github.com/mark3labs/mcp-go/mcp"
1011
"github.com/stretchr/testify/require"
1112
"golang.org/x/net/context"
@@ -17,12 +18,12 @@ type McpClient struct {
1718
*client.Client
1819
}
1920

20-
func NewMcpClient(t *testing.T, mcpHttpServer http.Handler) *McpClient {
21+
func NewMcpClient(t *testing.T, mcpHttpServer http.Handler, options ...transport.StreamableHTTPCOption) *McpClient {
2122
require.NotNil(t, mcpHttpServer, "McpHttpServer must be provided")
2223
var err error
2324
ret := &McpClient{ctx: t.Context()}
2425
ret.testServer = httptest.NewServer(mcpHttpServer)
25-
ret.Client, err = client.NewStreamableHttpClient(ret.testServer.URL + "/mcp")
26+
ret.Client, err = client.NewStreamableHttpClient(ret.testServer.URL+"/mcp", options...)
2627
require.NoError(t, err, "Expected no error creating MCP client")
2728
err = ret.Start(t.Context())
2829
require.NoError(t, err, "Expected no error starting MCP client")

‎pkg/mcp/common_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -443,9 +443,9 @@ func (s *BaseMcpSuite) TearDownTest() {
443443
}
444444
}
445445

446-
func (s *BaseMcpSuite) InitMcpClient() {
446+
func (s *BaseMcpSuite) InitMcpClient(options ...transport.StreamableHTTPCOption) {
447447
var err error
448448
s.mcpServer, err = NewServer(Configuration{StaticConfig: s.Cfg})
449449
s.Require().NoError(err, "Expected no error creating MCP server")
450-
s.McpClient = test.NewMcpClient(s.T(), s.mcpServer.ServeHTTP(nil))
450+
s.McpClient = test.NewMcpClient(s.T(), s.mcpServer.ServeHTTP(nil), options...)
451451
}

‎pkg/mcp/mcp_test.go‎

Lines changed: 50 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,9 @@ import (
1010
"time"
1111

1212
"github.com/containers/kubernetes-mcp-server/internal/test"
13-
"github.com/mark3labs/mcp-go/client"
13+
"github.com/mark3labs/mcp-go/client/transport"
1414
"github.com/mark3labs/mcp-go/mcp"
15+
"github.com/stretchr/testify/suite"
1516
)
1617

1718
func TestWatchKubeConfig(t *testing.T) {
@@ -48,16 +49,19 @@ func TestWatchKubeConfig(t *testing.T) {
4849
})
4950
}
5051

51-
func TestSseHeaders(t *testing.T) {
52-
mockServer := test.NewMockServer()
53-
defer mockServer.Close()
54-
before := func(c *mcpContext) {
55-
c.withKubeConfig(mockServer.Config())
56-
c.clientOptions = append(c.clientOptions, client.WithHeaders(map[string]string{"kubernetes-authorization": "Bearer a-token-from-mcp-client"}))
57-
}
58-
pathHeaders := make(map[string]http.Header, 0)
59-
mockServer.Handle(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
60-
pathHeaders[req.URL.Path] = req.Header.Clone()
52+
type McpHeadersSuite struct {
53+
BaseMcpSuite
54+
mockServer *test.MockServer
55+
pathHeaders map[string]http.Header
56+
}
57+
58+
func (s *McpHeadersSuite) SetupTest() {
59+
s.BaseMcpSuite.SetupTest()
60+
s.mockServer = test.NewMockServer()
61+
s.Cfg.KubeConfig = s.mockServer.KubeconfigFile(s.T())
62+
s.pathHeaders = make(map[string]http.Header)
63+
s.mockServer.Handle(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
64+
s.pathHeaders[req.URL.Path] = req.Header.Clone()
6165
// Request Performed by DiscoveryClient to Kube API (Get API Groups legacy -core-)
6266
if req.URL.Path == "/api" {
6367
w.Header().Set("Content-Type", "application/json")
@@ -90,38 +94,42 @@ func TestSseHeaders(t *testing.T) {
9094
}
9195
w.WriteHeader(404)
9296
}))
93-
testCaseWithContext(t, &mcpContext{before: before}, func(c *mcpContext) {
94-
_, _ = c.callTool("pods_list", map[string]interface{}{})
95-
t.Run("DiscoveryClient propagates headers to Kube API", func(t *testing.T) {
96-
if len(pathHeaders) == 0 {
97-
t.Fatalf("No requests were made to Kube API")
98-
}
99-
if pathHeaders["/api"] == nil || pathHeaders["/api"].Get("Authorization") != "Bearer a-token-from-mcp-client" {
100-
t.Fatalf("Overridden header Authorization not found in request to /api")
101-
}
102-
if pathHeaders["/apis"] == nil || pathHeaders["/apis"].Get("Authorization") != "Bearer a-token-from-mcp-client" {
103-
t.Fatalf("Overridden header Authorization not found in request to /apis")
104-
}
105-
if pathHeaders["/api/v1"] == nil || pathHeaders["/api/v1"].Get("Authorization") != "Bearer a-token-from-mcp-client" {
106-
t.Fatalf("Overridden header Authorization not found in request to /api/v1")
107-
}
97+
}
98+
99+
func (s *McpHeadersSuite) TearDownTest() {
100+
s.BaseMcpSuite.TearDownTest()
101+
if s.mockServer != nil {
102+
s.mockServer.Close()
103+
}
104+
}
105+
106+
func (s *McpHeadersSuite) TestAuthorizationHeaderPropagation() {
107+
cases := []string{"kubernetes-authorization", "Authorization"}
108+
for _, header := range cases {
109+
s.InitMcpClient(transport.WithHTTPHeaders(map[string]string{header: "Bearer a-token-from-mcp-client"}))
110+
_, _ = s.CallTool("pods_list", map[string]interface{}{})
111+
s.Require().Greater(len(s.pathHeaders), 0, "No requests were made to Kube API")
112+
s.Run("DiscoveryClient propagates "+header+" header to Kube API", func() {
113+
s.Require().NotNil(s.pathHeaders["/api"], "No requests were made to /api")
114+
s.Equal("Bearer a-token-from-mcp-client", s.pathHeaders["/api"].Get("Authorization"), "Overridden header Authorization not found in request to /api")
115+
s.Require().NotNil(s.pathHeaders["/apis"], "No requests were made to /apis")
116+
s.Equal("Bearer a-token-from-mcp-client", s.pathHeaders["/apis"].Get("Authorization"), "Overridden header Authorization not found in request to /apis")
117+
s.Require().NotNil(s.pathHeaders["/api/v1"], "No requests were made to /api/v1")
118+
s.Equal("Bearer a-token-from-mcp-client", s.pathHeaders["/api/v1"].Get("Authorization"), "Overridden header Authorization not found in request to /api/v1")
108119
})
109-
t.Run("DynamicClient propagates headers to Kube API", func(t *testing.T) {
110-
if len(pathHeaders) == 0 {
111-
t.Fatalf("No requests were made to Kube API")
112-
}
113-
if pathHeaders["/api/v1/namespaces/default/pods"] == nil || pathHeaders["/api/v1/namespaces/default/pods"].Get("Authorization") != "Bearer a-token-from-mcp-client" {
114-
t.Fatalf("Overridden header Authorization not found in request to /api/v1/namespaces/default/pods")
115-
}
120+
s.Run("DynamicClient propagates "+header+" header to Kube API", func() {
121+
s.Require().NotNil(s.pathHeaders["/api/v1/namespaces/default/pods"], "No requests were made to /api/v1/namespaces/default/pods")
122+
s.Equal("Bearer a-token-from-mcp-client", s.pathHeaders["/api/v1/namespaces/default/pods"].Get("Authorization"), "Overridden header Authorization not found in request to /api/v1/namespaces/default/pods")
116123
})
117-
_, _ = c.callTool("pods_delete", map[string]interface{}{"name": "a-pod-to-delete"})
118-
t.Run("kubernetes.Interface propagates headers to Kube API", func(t *testing.T) {
119-
if len(pathHeaders) == 0 {
120-
t.Fatalf("No requests were made to Kube API")
121-
}
122-
if pathHeaders["/api/v1/namespaces/default/pods/a-pod-to-delete"] == nil || pathHeaders["/api/v1/namespaces/default/pods/a-pod-to-delete"].Get("Authorization") != "Bearer a-token-from-mcp-client" {
123-
t.Fatalf("Overridden header Authorization not found in request to /api/v1/namespaces/default/pods/a-pod-to-delete")
124-
}
124+
_, _ = s.CallTool("pods_delete", map[string]interface{}{"name": "a-pod-to-delete"})
125+
s.Run("kubernetes.Interface propagates "+header+" header to Kube API", func() {
126+
s.Require().NotNil(s.pathHeaders["/api/v1/namespaces/default/pods/a-pod-to-delete"], "No requests were made to /api/v1/namespaces/default/pods/a-pod-to-delete")
127+
s.Equal("Bearer a-token-from-mcp-client", s.pathHeaders["/api/v1/namespaces/default/pods/a-pod-to-delete"].Get("Authorization"), "Overridden header Authorization not found in request to /api/v1/namespaces/default/pods/a-pod-to-delete")
125128
})
126-
})
129+
130+
}
131+
}
132+
133+
func TestMcpHeaders(t *testing.T) {
134+
suite.Run(t, new(McpHeadersSuite))
127135
}

0 commit comments

Comments
 (0)