Merge pull request #6404 from jungaretti/jungaretti/add-rebuild-command

Add cs rebuild command

internal/codespaces/api/api.go

@@ -210,6 +210,8 @@ const (
 	CodespaceStateShutdown = "Shutdown"
 	// CodespaceStateStarting is the state for a starting codespace environment.
 	CodespaceStateStarting = "Starting"
+	// CodespaceStateRebuilding is the state for a rebuilding codespace environment.
+	CodespaceStateRebuilding = "Rebuilding"
 )
 
 type CodespaceConnection struct {

pkg/cmd/codespace/common.go

@@ -66,6 +66,7 @@ type liveshareSession interface {
 	StartSharing(context.Context, string, int) (liveshare.ChannelID, error)
 	StartSSHServer(context.Context) (int, string, error)
 	StartSSHServerWithOptions(context.Context, liveshare.StartSSHServerOptions) (int, string, error)
+	RebuildContainer(context.Context) error
 }
 
 // Connects to a codespace using Live Share and returns that session

pkg/cmd/codespace/rebuild.go

@@ -0,0 +1,56 @@
+package codespace
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cli/cli/v2/internal/codespaces/api"
+	"github.com/spf13/cobra"
+)
+
+func newRebuildCmd(app *App) *cobra.Command {
+	var codespace string
+
+	rebuildCmd := &cobra.Command{
+		Use:   "rebuild",
+		Short: "Rebuild a codespace",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return app.Rebuild(cmd.Context(), codespace)
+		},
+	}
+
+	rebuildCmd.Flags().StringVarP(&codespace, "codespace", "c", "", "Name of the codespace")
+
+	return rebuildCmd
+}
+
+func (a *App) Rebuild(ctx context.Context, codespaceName string) (err error) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	codespace, err := getOrChooseCodespace(ctx, a.apiClient, codespaceName)
+	if err != nil {
+		return err
+	}
+
+	// There's no need to rebuild again because users can't modify their codespace while it rebuilds
+	if codespace.State == api.CodespaceStateRebuilding {
+		fmt.Fprintf(a.io.Out, "%s is already rebuilding\n", codespace.Name)
+		return nil
+	}
+
+	session, err := startLiveShareSession(ctx, codespace, a, false, "")
+	if err != nil {
+		return fmt.Errorf("starting Live Share session: %w", err)
+	}
+	defer safeClose(session, &err)
+
+	err = session.RebuildContainer(ctx)
+	if err != nil {
+		return fmt.Errorf("rebuilding codespace via session: %w", err)
+	}
+
+	fmt.Fprintf(a.io.Out, "%s is rebuilding\n", codespace.Name)
+	return nil
+}

pkg/cmd/codespace/rebuild_test.go

@@ -0,0 +1,36 @@
+package codespace
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cli/cli/v2/internal/codespaces/api"
+	"github.com/cli/cli/v2/pkg/iostreams"
+)
+
+func TestAlreadyRebuildingCodespace(t *testing.T) {
+	rebuildingCodespace := &api.Codespace{
+		Name:  "rebuildingCodespace",
+		State: api.CodespaceStateRebuilding,
+	}
+	app := testingRebuildApp(*rebuildingCodespace)
+
+	err := app.Rebuild(context.Background(), "rebuildingCodespace")
+	if err != nil {
+		t.Errorf("rebuilding a codespace that was already rebuilding: %v", err)
+	}
+}
+
+func testingRebuildApp(mockCodespace api.Codespace) *App {
+	apiMock := &apiClientMock{
+		GetCodespaceFunc: func(_ context.Context, name string, _ bool) (*api.Codespace, error) {
+			if name == mockCodespace.Name {
+				return &mockCodespace, nil
+			}
+			return nil, nil
+		},
+	}
+
+	ios, _, _, _ := iostreams.Test()
+	return NewApp(ios, nil, apiMock, nil)
+}

pkg/cmd/codespace/root.go

@@ -22,6 +22,7 @@ func NewRootCmd(app *App) *cobra.Command {
 	root.AddCommand(newCpCmd(app))
 	root.AddCommand(newStopCmd(app))
 	root.AddCommand(newSelectCmd(app))
+	root.AddCommand(newRebuildCmd(app))
 
 	return root
 }

pkg/liveshare/session.go

@@ -123,6 +123,20 @@ func (s *Session) StartJupyterServer(ctx context.Context) (int, string, error) {
 	return port, response.ServerUrl, nil
 }
 
+func (s *Session) RebuildContainer(ctx context.Context) error {
+	var rebuildSuccess bool
+	err := s.rpc.do(ctx, "IEnvironmentConfigurationService.rebuildContainer", nil, &rebuildSuccess)
+	if err != nil {
+		return fmt.Errorf("invoking rebuild RPC: %w", err)
+	}
+
+	if !rebuildSuccess {
+		return fmt.Errorf("couldn't rebuild codespace")
+	}
+
+	return nil
+}
+
 // heartbeat runs until context cancellation, periodically checking whether there is a
 // reason to keep the connection alive, and if so, notifying the Live Share host to do so.
 // Heartbeat ensures it does not send more than one request every "interval" to ratelimit

pkg/liveshare/session_test.go

@@ -399,6 +399,30 @@ func TestSessionHeartbeat(t *testing.T) {
 	}
 }
 
+func TestRebuild(t *testing.T) {
+	requestCount := 0
+	getSharedServers := func(conn *jsonrpc2.Conn, req *jsonrpc2.Request) (interface{}, error) {
+		requestCount++
+		return true, nil
+	}
+	testServer, session, err := makeMockSession(
+		livesharetest.WithService("IEnvironmentConfigurationService.rebuildContainer", getSharedServers),
+	)
+	if err != nil {
+		t.Fatalf("creating mock session: %v", err)
+	}
+	defer testServer.Close()
+
+	err = session.RebuildContainer(context.Background())
+	if err != nil {
+		t.Fatalf("rebuilding codespace via mock session: %v", err)
+	}
+
+	if requestCount == 0 {
+		t.Fatalf("no requests were made")
+	}
+}
+
 type mockLogger struct {
 	sync.Mutex
 	buf *bytes.Buffer