From bf81090458b5cc1d80733cabd57df6b5df44b6f7 Mon Sep 17 00:00:00 2001
From: Juan Antonio Osorio <ozz@stacklok.com>
Date: Fri, 31 Oct 2025 13:54:45 +0200
Subject: [PATCH] Fix flaky vMCP health endpoint tests
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The health endpoint tests were failing intermittently in CI due to port
conflicts when running parallel tests. All test instances tried to bind
to the default port 4483, causing "address already in use" errors.

Changes:
- Add Ready() channel to vMCP server to signal when listener is created
- Use networking.FindAvailable() to assign unique random ports to tests
- Improve test synchronization by waiting on Ready() channel
- Add error channel to catch and report startup failures

The tests now run reliably with 20 parallel iterations.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 pkg/vmcp/server/health_test.go | 36 ++++++++++++++++++++++++----------
 pkg/vmcp/server/server.go      | 17 ++++++++++++++++
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/pkg/vmcp/server/health_test.go b/pkg/vmcp/server/health_test.go
index 03660efc7..6b10c5d95 100644
--- a/pkg/vmcp/server/health_test.go
+++ b/pkg/vmcp/server/health_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
+	"github.com/stacklok/toolhive/pkg/networking"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
 	"github.com/stacklok/toolhive/pkg/vmcp/mocks"
@@ -30,11 +31,15 @@ func createTestServer(t *testing.T) *server.Server {
 	mockBackendClient := mocks.NewMockBackendClient(ctrl)
 	rt := router.NewDefaultRouter()
 
+	// Find an available port for parallel test execution
+	port := networking.FindAvailable()
+	require.NotZero(t, port, "Failed to find available port")
+
 	srv := server.New(&server.Config{
 		Name:    "test-vmcp",
 		Version: "1.0.0",
 		Host:    "127.0.0.1",
-		Port:    0, // Random port for parallel tests
+		Port:    port,
 	}, rt, mockBackendClient)
 
 	// Register minimal capabilities
@@ -54,17 +59,28 @@ func createTestServer(t *testing.T) *server.Server {
 	require.NoError(t, err)
 
 	// Start server in background
-	ctx, cancel := context.WithCancel(context.Background())
-	go func() { _ = srv.Start(ctx) }()
+	ctx, cancel := context.WithCancel(t.Context())
+	t.Cleanup(cancel)
 
-	// Wait for server to start
-	time.Sleep(100 * time.Millisecond)
+	errCh := make(chan error, 1)
+	go func() {
+		if err := srv.Start(ctx); err != nil {
+			errCh <- err
+		}
+	}()
+
+	// Wait for server to be ready (with timeout)
+	select {
+	case <-srv.Ready():
+		// Server is ready to accept connections
+	case err := <-errCh:
+		t.Fatalf("Server failed to start: %v", err)
+	case <-time.After(5 * time.Second):
+		t.Fatalf("Server did not become ready within 5s (address: %s)", srv.Address())
+	}
 
-	// Cleanup
-	t.Cleanup(func() {
-		cancel()
-		time.Sleep(50 * time.Millisecond)
-	})
+	// Give the HTTP server a moment to start accepting connections
+	time.Sleep(10 * time.Millisecond)
 
 	return srv
 }
diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go
index 190c72f0a..28dffcd0e 100644
--- a/pkg/vmcp/server/server.go
+++ b/pkg/vmcp/server/server.go
@@ -103,6 +103,11 @@ type Server struct {
 	// The mark3labs SDK calls our sessionIDAdapter, which delegates to this manager.
 	// The SDK does NOT manage sessions itself - it only provides the interface.
 	sessionManager *session.Manager
+
+	// Ready channel signals when the server is ready to accept connections.
+	// Closed once the listener is created and serving.
+	ready     chan struct{}
+	readyOnce sync.Once
 }
 
 // New creates a new Virtual MCP Server instance.
@@ -148,6 +153,7 @@ func New(
 		router:         rt,
 		backendClient:  backendClient,
 		sessionManager: sessionManager,
+		ready:          make(chan struct{}),
 	}
 }
 
@@ -262,6 +268,11 @@ func (s *Server) Start(ctx context.Context) error {
 		}
 	}()
 
+	// Signal that the server is ready (listener created and serving started)
+	s.readyOnce.Do(func() {
+		close(s.ready)
+	})
+
 	// Wait for either context cancellation or server error
 	select {
 	case <-ctx.Done():
@@ -600,3 +611,9 @@ func (*Server) handleHealth(w http.ResponseWriter, _ *http.Request) {
 func (s *Server) SessionManager() *session.Manager {
 	return s.sessionManager
 }
+
+// Ready returns a channel that is closed when the server is ready to accept connections.
+// This is useful for testing and synchronization.
+func (s *Server) Ready() <-chan struct{} {
+	return s.ready
+}