Tests for update port visibility

pkg/cmd/codespace/ports.go

@@ -260,6 +260,7 @@ func (a *App) UpdatePortVisibility(ctx context.Context, codespaceName string, ar
 		ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
 		defer cancel()
 
+		fmt.Println("watiing for update")
 		if err := a.waitForPortUpdate(ctx, session, port.number); err != nil {
 			return fmt.Errorf("error waiting for port update: %w", err)
 		}
@@ -291,25 +292,26 @@ func (a *App) waitForPortUpdate(ctx context.Context, session *liveshare.Session,
 	failure := session.WaitForEvent("sharingFailed")
 
 	for {
+		var pd portData
 		select {
 		case <-ctx.Done():
 			return fmt.Errorf("timeout waiting for server sharing to succeed or fail")
 		case b := <-success:
-			if err := json.Unmarshal(b, &portData); err != nil {
+			if err := json.Unmarshal(b, &pd); err != nil {
 				return fmt.Errorf("error unmarshaling port data: %w", err)
 			}
-			if portData.Port == port && portData.ChangeKind == portChangeKindUpdate {
+			if pd.Port == port && pd.ChangeKind == portChangeKindUpdate {
 				return nil
 			}
 		case b := <-failure:
-			if err := json.Unmarshal(b, &portData); err != nil {
+			if err := json.Unmarshal(b, &pd); err != nil {
 				return fmt.Errorf("error unmarshaling port data: %w", err)
 			}
-			if portData.Port == port && portData.ChangeKind == portChangeKindUpdate {
-				if portData.StatusCode == http.StatusForbidden {
+			if pd.Port == port && pd.ChangeKind == portChangeKindUpdate {
+				if pd.StatusCode == http.StatusForbidden {
 					return errors.New("organization admin has forbidden this privacy setting")
 				}
-				return errors.New(portData.ErrorDetail)
+				return errors.New(pd.ErrorDetail)
 			}
 		}
 	}

pkg/cmd/codespace/ports_test.go

@@ -0,0 +1,97 @@
+package codespace
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"testing"
+
+	"github.com/cli/cli/v2/internal/codespaces/api"
+	"github.com/cli/cli/v2/pkg/iostreams"
+	livesharetest "github.com/cli/cli/v2/pkg/liveshare/test"
+	"github.com/sourcegraph/jsonrpc2"
+)
+
+type joinWorkspaceResult struct {
+	SessionNumber int `json:"sessionNumber"`
+}
+
+func TestPortsUpdateVisibility(t *testing.T) {
+	joinWorkspace := func(req *jsonrpc2.Request) (interface{}, error) {
+		return joinWorkspaceResult{1}, nil
+	}
+	const sessionToken = "session-token"
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ch := make(chan float64, 1)
+	updateSharedVisibility := func(rpcReq *jsonrpc2.Request) (interface{}, error) {
+		var req []interface{}
+		if err := json.Unmarshal(*rpcReq.Params, &req); err != nil {
+			return nil, fmt.Errorf("unmarshal req: %w", err)
+		}
+
+		ch <- req[0].(float64)
+		return nil, nil
+	}
+	testServer, err := livesharetest.NewServer(
+		livesharetest.WithNonSecure(),
+		livesharetest.WithPassword(sessionToken),
+		livesharetest.WithService("workspace.joinWorkspace", joinWorkspace),
+		livesharetest.WithService("serverSharing.updateSharedServerPrivacy", updateSharedVisibility),
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type rpcMessage struct {
+		Method string
+		Params portData
+	}
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case port := <-ch:
+				testServer.WriteToObjectStream(rpcMessage{
+					Method: "sharingSucceeded",
+					Params: portData{
+						Port:       int(port),
+						ChangeKind: portChangeKindUpdate,
+					},
+				})
+			}
+		}
+	}()
+
+	mockApi := &apiClientMock{
+		GetCodespaceFunc: func(ctx context.Context, codespaceName string, includeConnection bool) (*api.Codespace, error) {
+			return &api.Codespace{
+				Name:  "codespace-name",
+				State: api.CodespaceStateAvailable,
+				Connection: api.CodespaceConnection{
+					SessionID:      "session-id",
+					SessionToken:   sessionToken,
+					RelayEndpoint:  testServer.URL(),
+					RelaySAS:       "relay-sas",
+					HostPublicKeys: []string{livesharetest.SSHPublicKey},
+				},
+			}, nil
+		},
+	}
+
+	fmt.Println(testServer)
+	io, _, _, _ := iostreams.Test()
+	a := &App{
+		io:        io,
+		apiClient: mockApi,
+	}
+
+	err = a.UpdatePortVisibility(ctx, "codespace-name", []string{"80:80", "9999:9999"})
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+}

pkg/liveshare/client.go

@@ -57,7 +57,13 @@ func (opts *Options) uri(action string) (string, error) {
 
 	sas := url.QueryEscape(opts.RelaySAS)
 	uri := opts.RelayEndpoint
-	uri = strings.Replace(uri, "sb:", "wss:", -1)
+
+	if strings.HasPrefix(uri, "http:") {
+		uri = strings.Replace(uri, "http:", "ws:", 1)
+	} else {
+		uri = strings.Replace(uri, "sb:", "wss:", -1)
+	}
+
 	uri = strings.Replace(uri, ".net/", ".net:443/$hc/", 1)
 	uri = uri + "?sb-hc-action=" + action + "&sb-hc-token=" + sas
 	return uri, nil

pkg/liveshare/test/server.go

@@ -42,6 +42,9 @@ type Server struct {
 	sshConfig      *ssh.ServerConfig
 	httptestServer *httptest.Server
 	errCh          chan error
+	nonSecure      bool
+
+	objectStream jsonrpc2.ObjectStream
 }
 
 // NewServer creates a new Server. ServerOptions can be passed to configure
@@ -65,7 +68,12 @@ func NewServer(opts ...ServerOption) (*Server, error) {
 	server.sshConfig.AddHostKey(privateKey)
 
 	server.errCh = make(chan error, 1)
-	server.httptestServer = httptest.NewTLSServer(http.HandlerFunc(makeConnection(server)))
+
+	if server.nonSecure {
+		server.httptestServer = httptest.NewServer(http.HandlerFunc(makeConnection(server)))
+	} else {
+		server.httptestServer = httptest.NewTLSServer(http.HandlerFunc(makeConnection(server)))
+	}
 	return server, nil
 }
 
@@ -80,6 +88,13 @@ func WithPassword(password string) ServerOption {
 	}
 }
 
+func WithNonSecure() ServerOption {
+	return func(s *Server) error {
+		s.nonSecure = true
+		return nil
+	}
+}
+
 // WithService accepts a mock RPC service for the Server to invoke.
 func WithService(serviceName string, handler RPCHandleFunc) ServerOption {
 	return func(s *Server) error {
@@ -134,6 +149,13 @@ func (s *Server) Err() <-chan error {
 	return s.errCh
 }
 
+func (s *Server) WriteToObjectStream(obj interface{}) error {
+	if s.objectStream == nil {
+		return errors.New("object stream not set")
+	}
+	return s.objectStream.WriteObject(obj)
+}
+
 var upgrader = websocket.Upgrader{}
 
 func makeConnection(server *Server) http.HandlerFunc {
@@ -300,6 +322,8 @@ func forwardStream(ctx context.Context, server *Server, streamName string, chann
 
 func handleChannel(server *Server, channel ssh.Channel) {
 	stream := jsonrpc2.NewBufferedStream(channel, jsonrpc2.VSCodeObjectCodec{})
+	server.objectStream = stream
+
 	jsonrpc2.NewConn(context.Background(), stream, newRPCHandler(server))
 }