cmbrose · Jan 17, 2025 · Jan 8, 2025 · Jan 10, 2025 · Jan 13, 2025
@@ -240,6 +240,8 @@ const (
 	CodespaceStateAvailable = "Available"
 	// CodespaceStateShutdown is the state for a shutdown codespace environment.
 	CodespaceStateShutdown = "Shutdown"
+	// CodespaceStateShuttingDown is the state for a shutting down codespace environment.
+	CodespaceStateShuttingDown = "ShuttingDown"
 	// CodespaceStateStarting is the state for a starting codespace environment.
 	CodespaceStateStarting = "Starting"
 	// CodespaceStateRebuilding is the state for a rebuilding codespace environment.

@@ -13,6 +13,15 @@ import (
 	"github.com/cli/cli/v2/internal/codespaces/connection"
 )
 
+// codespaceStatePollingBackoff is the delay between state polls while waiting for codespaces to become
+// available. It's only exposed so that it can be shortened for testing, otherwise it should not be changed
+var codespaceStatePollingBackoff backoff.BackOff = backoff.NewExponentialBackOff(
+	backoff.WithInitialInterval(1*time.Second),
+	backoff.WithMultiplier(1.02),
+	backoff.WithMaxInterval(10*time.Second),
+	backoff.WithMaxElapsedTime(5*time.Minute),
+)
+
 func connectionReady(codespace *api.Codespace) bool {
 	// If the codespace is not available, it is not ready
 	if codespace.State != api.CodespaceStateAvailable {
@@ -67,41 +76,53 @@ func GetCodespaceConnection(ctx context.Context, progress progressIndicator, api
 
 // waitUntilCodespaceConnectionReady waits for a Codespace to be running and is able to be connected to.
 func waitUntilCodespaceConnectionReady(ctx context.Context, progress progressIndicator, apiClient apiClient, codespace *api.Codespace) (*api.Codespace, error) {
-	if codespace.State != api.CodespaceStateAvailable {
-		progress.StartProgressIndicatorWithLabel("Starting codespace")
-		defer progress.StopProgressIndicator()
-		if err := apiClient.StartCodespace(ctx, codespace.Name); err != nil {
-			return nil, fmt.Errorf("error starting codespace: %w", err)
-		}
+	if connectionReady(codespace) {
+		return codespace, nil
 	}
 
-	if !connectionReady(codespace) {
-		expBackoff := backoff.NewExponentialBackOff()
-		expBackoff.Multiplier = 1.1
-		expBackoff.MaxInterval = 10 * time.Second
-		expBackoff.MaxElapsedTime = 5 * time.Minute
+	progress.StartProgressIndicatorWithLabel("Waiting for codespace to become ready")
+	defer progress.StopProgressIndicator()
+
+	lastState := ""
+	firstRetry := true
 
-		err := backoff.Retry(func() error {
-			var err error
+	err := backoff.Retry(func() error {
+		var err error
+		if firstRetry {
+			firstRetry = false
+		} else {
 			codespace, err = apiClient.GetCodespace(ctx, codespace.Name, true)
 			if err != nil {
 				return backoff.Permanent(fmt.Errorf("error getting codespace: %w", err))
 			}
+		}
 
-			if connectionReady(codespace) {
-				return nil
-			}
+		if connectionReady(codespace) {
+			return nil
+		}
 
-			return &TimeoutError{message: "codespace not ready yet"}
-		}, backoff.WithContext(expBackoff, ctx))
-		if err != nil {
-			var timeoutErr *TimeoutError
-			if errors.As(err, &timeoutErr) {
-				return nil, errors.New("timed out while waiting for the codespace to start")
+		// Only react to changes in the state (so that we don't try to start the codespace twice)
+		if codespace.State != lastState {
+			if codespace.State == api.CodespaceStateShutdown {
+				err = apiClient.StartCodespace(ctx, codespace.Name)
+				if err != nil {
+					return backoff.Permanent(fmt.Errorf("error starting codespace: %w", err))
+				}
 			}
+		}
+
+		lastState = codespace.State
+
+		return &TimeoutError{message: "codespace not ready yet"}
+	}, backoff.WithContext(codespaceStatePollingBackoff, ctx))
 
-			return nil, err
+	if err != nil {
+		var timeoutErr *TimeoutError
+		if errors.As(err, &timeoutErr) {
+			return nil, errors.New("timed out while waiting for the codespace to start")
 		}
+
+		return nil, err
 	}
 
 	return codespace, nil

@@ -0,0 +1,212 @@
+package codespaces
+
+import (
+	"context"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	"github.com/cli/cli/v2/internal/codespaces/api"
+)
+
+func init() {
+	// Set the backoff to 0 for testing so that they run quickly
+	codespaceStatePollingBackoff = backoff.NewConstantBackOff(time.Second * 0)
+}
+
+// This is just enough to trick `connectionReady`
+var readyCodespace = &api.Codespace{
+	State: api.CodespaceStateAvailable,
+	Connection: api.CodespaceConnection{
+		TunnelProperties: api.TunnelProperties{
+			ConnectAccessToken:     "test",
+			ManagePortsAccessToken: "test",
+			ServiceUri:             "test",
+			TunnelId:               "test",
+			ClusterId:              "test",
+			Domain:                 "test",
+		},
+	},
+}
+
+func TestWaitUntilCodespaceConnectionReady_WhenAlreadyReady(t *testing.T) {
+	t.Parallel()
+
+	apiClient := &mockApiClient{}
+	result, err := waitUntilCodespaceConnectionReady(context.Background(), &mockProgressIndicator{}, apiClient, readyCodespace)
+	if err != nil {
+		t.Fatalf("Expected nil error, but was %v", err)
+	}
+	if result.State != api.CodespaceStateAvailable {
+		t.Fatalf("Expected final state to be %s, but was %s", api.CodespaceStateAvailable, result.State)
+	}
+}
+
+func TestWaitUntilCodespaceConnectionReady_PollsApi(t *testing.T) {
+	t.Parallel()
+
+	apiClient := &mockApiClient{
+		onGetCodespace: func() (*api.Codespace, error) {
+			return readyCodespace, nil
+		},
+	}
+	result, err := waitUntilCodespaceConnectionReady(context.Background(), &mockProgressIndicator{}, apiClient, &api.Codespace{State: api.CodespaceStateStarting})
+
+	if err != nil {
+		t.Fatalf("Expected nil error, but was %v", err)
+	}
+	if result.State != api.CodespaceStateAvailable {
+		t.Fatalf("Expected final state to be %s, but was %s", api.CodespaceStateAvailable, result.State)
+	}
+}
+
+func TestWaitUntilCodespaceConnectionReady_StartsCodespace(t *testing.T) {
+	t.Parallel()
+
+	codespace := &api.Codespace{State: api.CodespaceStateShutdown}
+
+	apiClient := &mockApiClient{
+		onGetCodespace: func() (*api.Codespace, error) {
+			return codespace, nil
+		},
+		onStartCodespace: func() error {
+			*codespace = *readyCodespace
+			return nil
+		},
+	}
+	result, err := waitUntilCodespaceConnectionReady(context.Background(), &mockProgressIndicator{}, apiClient, codespace)
+	if err != nil {
+		t.Fatalf("Expected nil error, but was %v", err)
+	}
+	if result.State != api.CodespaceStateAvailable {
+		t.Fatalf("Expected final state to be %s, but was %s", api.CodespaceStateAvailable, result.State)
+	}
+}
+
+func TestWaitUntilCodespaceConnectionReady_PollsCodespaceUntilReady(t *testing.T) {
+	t.Parallel()
+
+	codespace := &api.Codespace{State: api.CodespaceStateShutdown}
+	hasPolled := false
+
+	apiClient := &mockApiClient{
+		onGetCodespace: func() (*api.Codespace, error) {
+			if hasPolled {
+				*codespace = *readyCodespace
+			}
+
+			hasPolled = true
+
+			return codespace, nil
+		},
+		onStartCodespace: func() error {
+			codespace.State = api.CodespaceStateStarting
+			return nil
+		},
+	}
+	result, err := waitUntilCodespaceConnectionReady(context.Background(), &mockProgressIndicator{}, apiClient, codespace)
+	if err != nil {
+		t.Fatalf("Expected nil error, but was %v", err)
+	}
+	if result.State != api.CodespaceStateAvailable {
+		t.Fatalf("Expected final state to be %s, but was %s", api.CodespaceStateAvailable, result.State)
+	}
+}
+
+func TestWaitUntilCodespaceConnectionReady_WaitsForShutdownBeforeStarting(t *testing.T) {
+	t.Parallel()
+
+	codespace := &api.Codespace{State: api.CodespaceStateShuttingDown}
+
+	apiClient := &mockApiClient{
+		onGetCodespace: func() (*api.Codespace, error) {
+			// Make sure that we poll at least once before going to shutdown
+			if codespace.State == api.CodespaceStateShuttingDown {
+				codespace.State = api.CodespaceStateShutdown
+			}
+			return codespace, nil
+		},
+		onStartCodespace: func() error {
+			if codespace.State != api.CodespaceStateShutdown {
+				t.Fatalf("Codespace started from non-shutdown state: %s", codespace.State)
+			}
+			*codespace = *readyCodespace
+			return nil
+		},
+	}
+	result, err := waitUntilCodespaceConnectionReady(context.Background(), &mockProgressIndicator{}, apiClient, codespace)
+	if err != nil {
+		t.Fatalf("Expected nil error, but was %v", err)
+	}
+	if result.State != api.CodespaceStateAvailable {
+		t.Fatalf("Expected final state to be %s, but was %s", api.CodespaceStateAvailable, result.State)
+	}
+}
+
+func TestUntilCodespaceConnectionReady_DoesntStartTwice(t *testing.T) {
+	t.Parallel()
+
+	codespace := &api.Codespace{State: api.CodespaceStateShutdown}
+	didStart := false
+	didPollAfterStart := false
+
+	apiClient := &mockApiClient{
+		onGetCodespace: func() (*api.Codespace, error) {
+			// Make sure that we are in shutdown state for one poll after starting to make sure we don't try to start again
+			if didPollAfterStart {
+				*codespace = *readyCodespace
+			}
+
+			if didStart {
+				didPollAfterStart = true
+			}
+
+			return codespace, nil
+		},
+		onStartCodespace: func() error {
+			if didStart {
+				t.Fatal("Should not start multiple times")
+			}
+			didStart = true
+			return nil
+		},
+	}
+	result, err := waitUntilCodespaceConnectionReady(context.Background(), &mockProgressIndicator{}, apiClient, codespace)
+	if err != nil {
+		t.Fatalf("Expected nil error, but was %v", err)
+	}
+	if result.State != api.CodespaceStateAvailable {
+		t.Fatalf("Expected final state to be %s, but was %s", api.CodespaceStateAvailable, result.State)
+	}
+}
+
+type mockApiClient struct {
+	onStartCodespace func() error
+	onGetCodespace   func() (*api.Codespace, error)
+}
+
+func (m *mockApiClient) StartCodespace(ctx context.Context, name string) error {
+	if m.onStartCodespace == nil {
+		panic("onStartCodespace not set and StartCodespace was called")
+	}
+
+	return m.onStartCodespace()
+}
+
+func (m *mockApiClient) GetCodespace(ctx context.Context, name string, includeConnection bool) (*api.Codespace, error) {
+	if m.onGetCodespace == nil {
+		panic("onGetCodespace not set and GetCodespace was called")
+	}
+
+	return m.onGetCodespace()
+}
+
+func (m *mockApiClient) HTTPClient() (*http.Client, error) {
+	panic("Not implemented")
+}
+
+type mockProgressIndicator struct{}
+
+func (m *mockProgressIndicator) StartProgressIndicatorWithLabel(s string) {}
+func (m *mockProgressIndicator) StopProgressIndicator()                   {}