Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 26 additions & 10 deletions pkg/vmcp/server/health_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand All @@ -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
}
Expand Down
17 changes: 17 additions & 0 deletions pkg/vmcp/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -148,6 +153,7 @@ func New(
router: rt,
backendClient: backendClient,
sessionManager: sessionManager,
ready: make(chan struct{}),
}
}

Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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
}
Loading