Skip to content

Commit 27e0e81

Browse files
manusamanusa
authored
test(http): try to reduce test flakiness (containers#489)
Signed-off-by: Marc Nuri <[email protected]>
1 parent 74fd934 commit 27e0e81

File tree

5 files changed

+46
-12
lines changed

5 files changed

+46
-12
lines changed

‎internal/test/mcp.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@ import (
1313
)
1414

1515
func McpInitRequest() mcp.InitializeRequest {
16-
initRequest := mcp.InitializeRequest{}
16+
initRequest := mcp.InitializeRequest{
17+
Request: mcp.Request{Method: "initialize"},
18+
}
1719
initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
1820
initRequest.Params.ClientInfo = mcp.Implementation{Name: "test", Version: "1.33.7"}
1921
return initRequest

‎internal/test/test.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package test
33
import (
44
"fmt"
55
"net"
6+
"net/http"
67
"os"
78
"path/filepath"
89
"runtime"
@@ -49,3 +50,24 @@ func WaitForServer(tcpAddr *net.TCPAddr) error {
4950
}
5051
return err
5152
}
53+
54+
// WaitForHealthz waits for the /healthz endpoint to return a non-404 response
55+
func WaitForHealthz(tcpAddr *net.TCPAddr) error {
56+
url := fmt.Sprintf("http://%s/healthz", tcpAddr.String())
57+
var resp *http.Response
58+
var err error
59+
for i := 0; i < 100; i++ {
60+
resp, err = http.Get(url)
61+
if err == nil {
62+
_ = resp.Body.Close()
63+
if resp.StatusCode != http.StatusNotFound {
64+
return nil
65+
}
66+
}
67+
time.Sleep(50 * time.Millisecond)
68+
}
69+
if err != nil {
70+
return err
71+
}
72+
return fmt.Errorf("healthz endpoint returned 404 after retries")
73+
}

‎pkg/http/http_authorization_test.go‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ func (s *AuthorizationSuite) SetupTest() {
3131
s.BaseHttpSuite.SetupTest()
3232

3333
// Capture logs
34+
s.logBuffer.Reset()
3435
s.klogState = klog.CaptureState()
3536
flags := flag.NewFlagSet("test", flag.ContinueOnError)
3637
klog.InitFlags(flags)
@@ -59,14 +60,14 @@ func (s *AuthorizationSuite) TearDownTest() {
5960

6061
func (s *AuthorizationSuite) StartClient(options ...transport.StreamableHTTPCOption) {
6162
var err error
62-
s.mcpClient, err = client.NewStreamableHttpClient(fmt.Sprintf("http://127.0.0.1:%d/mcp", s.TcpAddr.Port), options...)
63+
s.mcpClient, err = client.NewStreamableHttpClient(fmt.Sprintf("http://127.0.0.1:%s/mcp", s.StaticConfig.Port), options...)
6364
s.Require().NoError(err, "Expected no error creating Streamable HTTP MCP client")
6465
err = s.mcpClient.Start(s.T().Context())
6566
s.Require().NoError(err, "Expected no error starting Streamable HTTP MCP client")
6667
}
6768

6869
func (s *AuthorizationSuite) HttpGet(authHeader string) *http.Response {
69-
req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d/mcp", s.TcpAddr.Port), nil)
70+
req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://127.0.0.1:%s/mcp", s.StaticConfig.Port), nil)
7071
s.Require().NoError(err, "Failed to create request")
7172
if authHeader != "" {
7273
req.Header.Set("Authorization", authHeader)
@@ -339,6 +340,7 @@ func (s *AuthorizationSuite) TestAuthorizationRawToken() {
339340
for _, c := range cases {
340341
s.StaticConfig.OAuthAudience = c.audience
341342
s.StaticConfig.ValidateToken = c.validateToken
343+
s.logBuffer.Reset()
342344
s.StartServer()
343345
s.StartClient(transport.WithHTTPHeaders(map[string]string{
344346
"Authorization": "Bearer " + tokenBasicNotExpired,
@@ -362,7 +364,9 @@ func (s *AuthorizationSuite) TestAuthorizationRawToken() {
362364
})
363365
})
364366
_ = s.mcpClient.Close()
367+
s.mcpClient = nil
365368
s.StopServer()
369+
s.Require().NoError(s.WaitForShutdown())
366370
}
367371
}
368372

@@ -407,7 +411,9 @@ func (s *AuthorizationSuite) TestAuthorizationOidcToken() {
407411
})
408412
})
409413
_ = s.mcpClient.Close()
414+
s.mcpClient = nil
410415
s.StopServer()
416+
s.Require().NoError(s.WaitForShutdown())
411417
}
412418
}
413419

@@ -440,6 +446,7 @@ func (s *AuthorizationSuite) TestAuthorizationOidcTokenExchange() {
440446
s.StaticConfig.StsClientSecret = "test-sts-client-secret"
441447
s.StaticConfig.StsAudience = "backend-audience"
442448
s.StaticConfig.StsScopes = []string{"backend-scope"}
449+
s.logBuffer.Reset()
443450
s.StartServer()
444451
s.StartClient(transport.WithHTTPHeaders(map[string]string{
445452
"Authorization": "Bearer " + validOidcClientToken,
@@ -463,7 +470,9 @@ func (s *AuthorizationSuite) TestAuthorizationOidcTokenExchange() {
463470
})
464471
})
465472
_ = s.mcpClient.Close()
473+
s.mcpClient = nil
466474
s.StopServer()
475+
s.Require().NoError(s.WaitForShutdown())
467476
}
468477
}
469478

‎pkg/http/http_mcp_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ func (s *McpTransportSuite) TearDownTest() {
2525
}
2626

2727
func (s *McpTransportSuite) TestSseTransport() {
28-
sseClient, sseClientErr := client.NewSSEMCPClient(fmt.Sprintf("http://127.0.0.1:%d/sse", s.TcpAddr.Port))
28+
sseClient, sseClientErr := client.NewSSEMCPClient(fmt.Sprintf("http://127.0.0.1:%s/sse", s.StaticConfig.Port))
2929
s.Require().NoError(sseClientErr, "Expected no error creating SSE MCP client")
3030
startErr := sseClient.Start(s.T().Context())
3131
s.Require().NoError(startErr, "Expected no error starting SSE MCP client")
@@ -44,7 +44,7 @@ func (s *McpTransportSuite) TestSseTransport() {
4444
}
4545

4646
func (s *McpTransportSuite) TestStreamableHttpTransport() {
47-
httpClient, httpClientErr := client.NewStreamableHttpClient(fmt.Sprintf("http://127.0.0.1:%d/mcp", s.TcpAddr.Port), transport.WithContinuousListening())
47+
httpClient, httpClientErr := client.NewStreamableHttpClient(fmt.Sprintf("http://127.0.0.1:%s/mcp", s.StaticConfig.Port), transport.WithContinuousListening())
4848
s.Require().NoError(httpClientErr, "Expected no error creating Streamable HTTP MCP client")
4949
startErr := httpClient.Start(s.T().Context())
5050
s.Require().NoError(startErr, "Expected no error starting Streamable HTTP MCP client")

‎pkg/http/http_test.go‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@ import (
3333
type BaseHttpSuite struct {
3434
suite.Suite
3535
MockServer *test.MockServer
36-
TcpAddr *net.TCPAddr
3736
StaticConfig *config.StaticConfig
3837
mcpServer *mcp.Server
3938
OidcProvider *oidc.Provider
@@ -43,18 +42,19 @@ type BaseHttpSuite struct {
4342
}
4443

4544
func (s *BaseHttpSuite) SetupTest() {
46-
var err error
4745
http.DefaultClient.Timeout = 10 * time.Second
4846
s.MockServer = test.NewMockServer()
49-
s.TcpAddr, err = test.RandomPortAddress()
50-
s.Require().NoError(err, "Expected no error getting random port address")
47+
s.MockServer.Handle(&test.DiscoveryClientHandler{})
5148
s.StaticConfig = config.Default()
5249
s.StaticConfig.KubeConfig = s.MockServer.KubeconfigFile(s.T())
53-
s.StaticConfig.Port = strconv.Itoa(s.TcpAddr.Port)
5450
}
5551

5652
func (s *BaseHttpSuite) StartServer() {
57-
var err error
53+
54+
tcpAddr, err := test.RandomPortAddress()
55+
s.Require().NoError(err, "Expected no error getting random port address")
56+
s.StaticConfig.Port = strconv.Itoa(tcpAddr.Port)
57+
5858
s.mcpServer, err = mcp.NewServer(mcp.Configuration{StaticConfig: s.StaticConfig})
5959
s.Require().NoError(err, "Expected no error creating MCP server")
6060
s.Require().NotNil(s.mcpServer, "MCP server should not be nil")
@@ -64,7 +64,8 @@ func (s *BaseHttpSuite) StartServer() {
6464
cancelCtx, s.StopServer = context.WithCancel(gc)
6565
group.Go(func() error { return Serve(cancelCtx, s.mcpServer, s.StaticConfig, s.OidcProvider, nil) })
6666
s.WaitForShutdown = group.Wait
67-
s.Require().NoError(test.WaitForServer(s.TcpAddr), "HTTP server did not start in time")
67+
s.Require().NoError(test.WaitForServer(tcpAddr), "HTTP server did not start in time")
68+
s.Require().NoError(test.WaitForHealthz(tcpAddr), "HTTP server /healthz endpoint did not respond with non-404 in time")
6869
}
6970

7071
func (s *BaseHttpSuite) TearDownTest() {

0 commit comments

Comments
 (0)