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 +}