Add grpc mock server + tests

internal/codespaces/grpc/client.go

@@ -41,7 +41,7 @@ type liveshareSession interface {
 	StartSharing(context.Context, string, int) (liveshare.ChannelID, error)
 }
 
-// Connects to the gRPC server on the given port
+// Finds a free port to listen on and creates a new gRPC client that connects to that port
 func Connect(ctx context.Context, session liveshareSession, token string) (*Client, error) {
 	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", 0))
 	if err != nil {
@@ -49,20 +49,34 @@ func Connect(ctx context.Context, session liveshareSession, token string) (*Clie
 	}
 
 	// Tunnel the remote gRPC server port to the local port
-	localGrpcServerPort := listener.Addr().(*net.TCPAddr).Port
+	localPort := listener.Addr().(*net.TCPAddr).Port
 	internalTunnelClosed := make(chan error, 1)
 	go func() {
 		fwd := liveshare.NewPortForwarder(session, codespacesInternalSessionName, codespacesInternalPort, true)
 		internalTunnelClosed <- fwd.ForwardToListener(ctx, listener)
 	}()
 
+	// Create the gRPC client
+	client, err := NewClient(ctx, session, token, localPort)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create client: %w", err)
+	}
+
+	// Attach the listener so we can close it later
+	client.listener = listener
+
+	return client, err
+}
+
+// Creates a new gRPC client that connects to the given port
+func NewClient(ctx context.Context, session liveshareSession, token string, localPort int) (*Client, error) {
 	// Attempt to connect to the given port
 	opts := []grpc.DialOption{
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithBlock(),
 	}
 	ctx, _ = context.WithTimeout(ctx, connectionTimeout)
-	conn, err := grpc.DialContext(ctx, fmt.Sprintf("127.0.0.1:%d", localGrpcServerPort), opts...)
+	conn, err := grpc.DialContext(ctx, fmt.Sprintf("127.0.0.1:%d", localPort), opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -70,7 +84,6 @@ func Connect(ctx context.Context, session liveshareSession, token string) (*Clie
 	g := &Client{
 		conn:          conn,
 		token:         token,
-		listener:      listener,
 		jupyterClient: jupyter.NewJupyterServerHostClient(conn),
 	}
 

internal/codespaces/grpc/client_test.go

@@ -0,0 +1,68 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"testing"
+
+	"github.com/cli/cli/v2/internal/codespaces/grpc/test"
+)
+
+func TestMain(m *testing.M) {
+	// Start the gRPC server in the background
+	go func() {
+		err := test.StartServer()
+		if err != nil {
+			panic(err)
+		}
+	}()
+
+	m.Run()
+}
+
+func connect(t *testing.T) (ctx context.Context, client *Client) {
+	ctx = context.Background()
+	client, err := NewClient(ctx, nil, "token", test.ServerPort)
+	client.listener = &net.TCPListener{} // mock listener so the close function doesn't panic
+	if err != nil {
+		t.Fatalf("error connecting to internal server: %v", err)
+	}
+
+	return ctx, client
+}
+
+// Test that the gRPC client returns the correct port and URL when the JupyterLab server starts successfully
+func TestStartJupyterServerSuccess(t *testing.T) {
+	ctx, client := connect(t)
+	defer client.Close()
+	port, url, err := client.StartJupyterServer(ctx)
+	if err != nil {
+		t.Fatalf("expected %v, got %v", nil, err)
+	}
+	if port != test.JupyterPort {
+		t.Fatalf("expected %d, got %d", test.JupyterPort, port)
+	}
+	if url != test.JupyterServerUrl {
+		t.Fatalf("expected %s, got %s", test.JupyterServerUrl, url)
+	}
+}
+
+// Test that the gRPC client returns an error when the JupyterLab server fails to start
+func TestStartJupyterServerFailure(t *testing.T) {
+	ctx, client := connect(t)
+	defer client.Close()
+	test.JupyterMessage = "error message"
+	test.JupyterResult = false
+	errorMessage := fmt.Sprintf("failed to start JupyterLab: %s", test.JupyterMessage)
+	port, url, err := client.StartJupyterServer(ctx)
+	if err.Error() != errorMessage {
+		t.Fatalf("expected %v, got %v", errorMessage, err)
+	}
+	if port != 0 {
+		t.Fatalf("expected %d, got %d", 0, port)
+	}
+	if url != "" {
+		t.Fatalf("expected %s, got %s", "", url)
+	}
+}

internal/codespaces/grpc/test/server.go

@@ -0,0 +1,52 @@
+package test
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"strconv"
+
+	"github.com/cli/cli/v2/internal/codespaces/grpc/jupyter"
+	"google.golang.org/grpc"
+)
+
+const (
+	ServerPort = 50051
+)
+
+var (
+	JupyterPort      = 1234
+	JupyterServerUrl = "http://localhost:1234?token=1234"
+	JupyterMessage   = ""
+	JupyterResult    = true
+)
+
+type server struct {
+	jupyter.UnimplementedJupyterServerHostServer
+}
+
+func (s *server) GetRunningServer(ctx context.Context, in *jupyter.GetRunningServerRequest) (*jupyter.GetRunningServerResponse, error) {
+	return &jupyter.GetRunningServerResponse{
+		Port:      strconv.Itoa(JupyterPort),
+		ServerUrl: JupyterServerUrl,
+		Message:   JupyterMessage,
+		Result:    JupyterResult,
+	}, nil
+}
+
+// Starts the mock gRPC server listening on port 50051
+func StartServer() error {
+	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", ServerPort))
+	if err != nil {
+		return fmt.Errorf("failed to listen: %v", err)
+	}
+	defer listener.Close()
+
+	s := grpc.NewServer()
+	jupyter.RegisterJupyterServerHostServer(s, &server{})
+	if err := s.Serve(listener); err != nil {
+		return fmt.Errorf("failed to serve: %v", err)
+	}
+
+	return nil
+}