[v18][flaky] TestReplaceHTTPResponse and TestHTTPReaderWriter to use synctest (#61308)

Author: greedy52

lib/utils/mcputils/http_test.go

@@ -1,3 +1,5 @@
+//go:build go1.24 && enablesynctest
+
 /*
  * Teleport
  * Copyright (C) 2025  Gravitational, Inc.
@@ -26,10 +28,11 @@ import (
 	"log/slog"
 	"maps"
 	"net/http"
+	"os"
 	"sync"
 	"sync/atomic"
 	"testing"
-	"time"
+	"testing/synctest"
 
 	"github.com/gravitational/trace"
 	mcpclient "github.com/mark3labs/mcp-go/client"
@@ -39,66 +42,86 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	listenerutils "github.com/gravitational/teleport/lib/utils/listener"
+	"github.com/gravitational/teleport/lib/utils/log/logtest"
 	"github.com/gravitational/teleport/lib/utils/mcptest"
 )
 
+func TestMain(m *testing.M) {
+	logtest.InitLogger(testing.Verbose)
+	os.Exit(m.Run())
+}
+
 func TestReplaceHTTPResponse(t *testing.T) {
-	t.Parallel()
-	ctx := t.Context()
-
-	// Set up a server.
-	mcpServer := mcptest.NewServerWithVersion("11.22.33")
-	httpServer := mcpserver.NewTestStreamableHTTPServer(mcpServer)
-	t.Cleanup(httpServer.Close)
-
-	// Set up a client with custom transport which calls "ReplaceHTTPResponse".
-	httpClientTransport := newTestReplaceHTTPResponseTransport()
-	mcpClientTransport, err := mcpclienttransport.NewStreamableHTTP(
-		httpServer.URL+"/mcp",
-		mcpclienttransport.WithHTTPBasicClient(
-			&http.Client{Transport: httpClientTransport},
-		),
-		mcpclienttransport.WithContinuousListening(),
-	)
-	require.NoError(t, err)
-	client := mcpclient.NewClient(mcpClientTransport)
-	require.NoError(t, client.Start(ctx))
-
-	// Initialize client and call a tool.
-	result := mcptest.MustInitializeClient(t, client)
-	require.Equal(t, "111.222.333", result.ServerInfo.Version)
-	mcptest.MustCallServerTool(t, client)
-	require.Equal(t, uint32(2), httpClientTransport.countMCPResponse.Load())
-
-	// Send notifications from server. Notifications will be sent through SSE.
-	require.EventuallyWithT(t, func(collect *assert.CollectT) {
-		assert.Greater(collect, httpClientTransport.getCountMethods()["GET"], 0)
-	}, 2*time.Second, 100*time.Millisecond, "client SSE connected")
-	mcpServer.SendNotificationToAllClients("notifications/test", nil)
-	mcpServer.SendNotificationToAllClients("notifications/test", nil)
-	require.EventuallyWithT(t, func(collect *assert.CollectT) {
-		assert.Equal(collect, uint32(2), httpClientTransport.countMCPNotification.Load())
-	}, 2*time.Second, 100*time.Millisecond, "expected to receive notification")
-
-	// Close client and count the requests.
-	require.NoError(t, client.Close())
-	require.Equal(t, map[string]int{
-		"GET":    1, // For listening on SSE events.
-		"POST":   3, // "initialize", "notifications/initialize", and "tools/call".
-		"DELETE": 1, // Close session.
-	}, httpClientTransport.getCountMethods())
+	synctest.Run(func() {
+		// TODO(greedy52) replace [synctest.Run + t.Run] with [synctest.Test]
+		// when switching to go1.25. The extra layer t.Run is a temp workaround
+		// to synctest deadlock with go 1.24 since t.Context() is used in many
+		// places for background goroutines to exit.
+		t.Run("synctest", func(t *testing.T) {
+			ctx := t.Context()
+
+			// Set up a server.
+			mcpServer := mcptest.NewServerWithVersion("11.22.33")
+			listener := makeHTTPServerWithInMemoryListener(t, mcpServer)
+
+			// Set up a client with custom transport which calls "ReplaceHTTPResponse".
+			httpClientTransport := newTestReplaceHTTPResponseTransport(listener)
+			mcpClientTransport, err := mcpclienttransport.NewStreamableHTTP(
+				"http://memory/mcp",
+				mcpclienttransport.WithHTTPBasicClient(
+					&http.Client{Transport: httpClientTransport},
+				),
+				mcpclienttransport.WithContinuousListening(),
+			)
+			require.NoError(t, err)
+			var countMCPNotification atomic.Uint32
+			client := mcpclient.NewClient(mcpClientTransport)
+			client.OnNotification(func(notification mcp.JSONRPCNotification) {
+				countMCPNotification.Add(1)
+			})
+			require.NoError(t, client.Start(ctx))
+
+			// Initialize client and call a tool.
+			result := mcptest.MustInitializeClient(t, client)
+			require.Equal(t, "111.222.333", result.ServerInfo.Version)
+			mcptest.MustCallServerTool(t, client)
+			require.Equal(t, uint32(2), httpClientTransport.countMCPResponse.Load())
+
+			// Send notifications from server. Notifications will be sent through SSE.
+			synctest.Wait() // Wait for client to establish the GET connection.
+			mcpServer.SendNotificationToAllClients("notifications/test", nil)
+			mcpServer.SendNotificationToAllClients("notifications/test", nil)
+			synctest.Wait() // Wait for client to receive notifications.
+			require.Equal(t, uint32(2), countMCPNotification.Load())
+
+			// Close client and count the requests.
+			require.NoError(t, client.Close())
+			synctest.Wait()
+			require.Equal(t, map[string]int{
+				"GET":    1, // For listening on SSE events.
+				"POST":   3, // "initialize", "notifications/initialize", and "tools/call".
+				"DELETE": 1, // Close session.
+			}, httpClientTransport.getCountMethods())
+		})
+	})
 }
 
 type testReplaceHTTPResponseTransport struct {
-	countMethods         map[string]int
-	countMethodsMu       sync.Mutex
-	countMCPResponse     atomic.Uint32
-	countMCPNotification atomic.Uint32
+	countMethods     map[string]int
+	countMethodsMu   sync.Mutex
+	countMCPResponse atomic.Uint32
+	client           http.Client
 }
 
-func newTestReplaceHTTPResponseTransport() *testReplaceHTTPResponseTransport {
+func newTestReplaceHTTPResponseTransport(inMemoryListener *listenerutils.InMemoryListener) *testReplaceHTTPResponseTransport {
 	return &testReplaceHTTPResponseTransport{
 		countMethods: make(map[string]int),
+		client: http.Client{
+			Transport: &http.Transport{
+				DialContext: inMemoryListener.DialContext,
+			},
+		},
 	}
 }
 
@@ -113,7 +136,7 @@ func (t *testReplaceHTTPResponseTransport) RoundTrip(r *http.Request) (*http.Res
 	t.countMethods[r.Method]++
 	t.countMethodsMu.Unlock()
 
-	resp, err := http.DefaultClient.Do(r)
+	resp, err := t.client.Do(r)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -147,55 +170,65 @@ func (t *testReplaceHTTPResponseTransport) ProcessResponse(_ context.Context, re
 }
 
 func (t *testReplaceHTTPResponseTransport) ProcessNotification(_ context.Context, notification *JSONRPCNotification) mcp.JSONRPCMessage {
-	t.countMCPNotification.Add(1)
 	return notification
 }
 
 func TestHTTPReaderWriter(t *testing.T) {
-	t.Parallel()
-	ctx := t.Context()
+	synctest.Run(func() {
+		// TODO(greedy52) replace [synctest.Run + t.Run] with [synctest.Test]
+		// when switching to go1.25. The extra layer t.Run is a temp workaround
+		// to synctest deadlock with go 1.24 since t.Context() is used in many
+		// places for background goroutines to exit.
+		t.Run("synctest", func(t *testing.T) {
+			ctx := t.Context()
 
-	// Set up an MCP server.
-	mcpServer := mcptest.NewServer()
-	httpServer := mcpserver.NewTestStreamableHTTPServer(mcpServer)
-	t.Cleanup(httpServer.Close)
+			// Set up an MCP server.
+			mcpServer := mcptest.NewServer()
+			listener := makeHTTPServerWithInMemoryListener(t, mcpServer)
 
-	// Create a proxy that converts from stdio to HTTP.
-	clientStdin, writeToClient := io.Pipe()
-	readFromClient, clientStdout := io.Pipe()
-	t.Cleanup(func() {
-		assert.NoError(t, trace.NewAggregate(
-			clientStdin.Close(), writeToClient.Close(),
-			readFromClient.Close(), clientStdout.Close(),
-		))
-	})
+			// Create a proxy that converts from stdio to HTTP.
+			clientStdin, writeToClient := io.Pipe()
+			readFromClient, clientStdout := io.Pipe()
+			t.Cleanup(func() {
+				assert.NoError(t, trace.NewAggregate(
+					clientStdin.Close(), writeToClient.Close(),
+					readFromClient.Close(), clientStdout.Close(),
+				))
+			})
 
-	serverReaderWriter, err := NewHTTPReaderWriter(ctx, httpServer.URL, mcpclienttransport.WithContinuousListening())
-	require.NoError(t, err)
-	defer serverReaderWriter.Close() // Send DELETE before server is shutdown
+			serverReaderWriter, err := NewHTTPReaderWriter(
+				ctx,
+				"http://memory/mcp",
+				mcpclienttransport.WithContinuousListening(),
+				mcpclienttransport.WithHTTPBasicClient(listener.MakeHTTPClient()),
+			)
+			require.NoError(t, err)
+			defer serverReaderWriter.Close() // Send DELETE before server is shutdown
+
+			clientTransportReader := NewStdioReader(readFromClient)
+			clientWriter := NewStdioMessageWriter(writeToClient)
+			proxyReaderWriter(t, clientTransportReader, clientWriter, serverReaderWriter, serverReaderWriter)
 
-	clientTransportReader := NewStdioReader(readFromClient)
-	clientWriter := NewStdioMessageWriter(writeToClient)
-	proxyReaderWriter(t, clientTransportReader, clientWriter, serverReaderWriter, serverReaderWriter)
+			// Make a "high-level" stdio MCP client and test the proxy.
+			var receivedNotifications []mcp.JSONRPCNotification
+			stdioClient := mcptest.NewStdioClient(t, clientStdin, clientStdout)
+			stdioClient.OnNotification(func(notification mcp.JSONRPCNotification) {
+				receivedNotifications = append(receivedNotifications, notification)
+			})
+			mcptest.MustInitializeClient(t, stdioClient)
+			mcptest.MustCallServerTool(t, stdioClient)
 
-	// Make a "high-level" stdio MCP client and test the proxy.
-	notificationsChan := make(chan mcp.JSONRPCNotification, 1)
-	stdioClient := mcptest.NewStdioClient(t, clientStdin, clientStdout)
-	stdioClient.OnNotification(func(notification mcp.JSONRPCNotification) {
-		notificationsChan <- notification
+			// Test listening notifications from server.
+			// First do a synctest.Wait until the client establish the listening
+			// stream before sending the notification from the server. Then do
+			// another synctest.Wait for the client to receive the notification.
+			synctest.Wait()
+			mcpServer.SendNotificationToAllClients("notifications/test", nil)
+			synctest.Wait()
+			require.Len(t, receivedNotifications, 1)
+			require.Equal(t, "notifications/test", receivedNotifications[0].Notification.Method)
+		})
 	})
-	mcptest.MustInitializeClient(t, stdioClient)
-	mcptest.MustCallServerTool(t, stdioClient)
-
-	// Test listening notifications from server.
-	mcpServer.SendNotificationToAllClients("notifications/test", nil)
-	select {
-	case notification := <-notificationsChan:
-		require.NotNil(t, notification)
-		require.Equal(t, "notifications/test", notification.Notification.Method)
-	case <-time.After(time.Second):
-		require.Fail(t, "timeout waiting for notification")
-	}
 }
 
 func proxyReaderWriter(
@@ -214,3 +247,18 @@ func proxyReaderWriter(
 	go clientMessageReader.Run(t.Context())
 	go serverMessageReader.Run(t.Context())
 }
+
+// makeHTTPServerWithInMemoryListener starts a streamable-HTTP MCP server using
+// an InMemoryListener. InMemoryListener is good to use with synctest.
+func makeHTTPServerWithInMemoryListener(t *testing.T, mcpServer *mcpserver.MCPServer) *listenerutils.InMemoryListener {
+	t.Helper()
+	listener := listenerutils.NewInMemoryListener()
+	httpServer := &http.Server{
+		Handler: mcpserver.NewStreamableHTTPServer(mcpServer),
+	}
+	go httpServer.Serve(listener)
+	t.Cleanup(func() {
+		httpServer.Close()
+	})
+	return listener
+}