From 48e3473a953b9502a840bf9ea40fb818dca05f5b Mon Sep 17 00:00:00 2001
From: Jose Garcia <jose.balius@gmail.com>
Date: Tue, 21 Sep 2021 18:18:30 -0400
Subject: [PATCH] PR Feedback

- Bring context.Timeout into the poller
- Accept duration and interval
- Other tidy up
---
 cmd/ghcs/create.go      | 31 ++++++++++++-------------------
 cmd/ghcs/create_test.go | 26 ++++++++------------------
 2 files changed, 20 insertions(+), 37 deletions(-)

diff --git a/cmd/ghcs/create.go b/cmd/ghcs/create.go
index fd54a170c..a13807094 100644
--- a/cmd/ghcs/create.go
+++ b/cmd/ghcs/create.go
@@ -84,19 +84,17 @@ func create(opts *createOptions) error {
 
 	log.Println("Creating your codespace...")
 
-	codespace, err := apiClient.CreateCodespace(
-		ctx, userResult.User, repository, machine, branch, locationResult.Location,
-	)
+	codespace, err := apiClient.CreateCodespace(ctx, userResult.User, repository, machine, branch, locationResult.Location)
 	if err != nil {
 		// This error is returned by the API when the initial creation fails with a retryable error.
 		// A retryable error means that GitHub will retry to re-create Codespace and clients should poll
 		// the API and attempt to fetch the Codespace for the next two minutes.
 		if err == api.ErrCreateAsyncRetry {
 			log.Print("Switching to async provisioning...")
-			pollctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
-			defer cancel()
 
-			codespace, err = pollForCodespace(pollctx, apiClient, log, userResult.User, codespace)
+			pollTimeout := 2 * time.Minute
+			pollInterval := 1 * time.Second
+			codespace, err = pollForCodespace(ctx, apiClient, log, pollTimeout, pollInterval, userResult.User.Login, codespace.Name)
 			log.Print("\n")
 
 			if err != nil {
@@ -125,13 +123,13 @@ type apiClient interface {
 	GetCodespace(context.Context, string, string, string) (*api.Codespace, error)
 }
 
-// pollForCodespace polls the Codespaces API every second fetching the codespace.
+// pollForCodespace polls the Codespaces GET endpoint on a given interval for a specified duration.
 // If it succeeds at fetching the codespace, we consider the codespace provisioned.
-// Context should be cancelled to stop polling.
-func pollForCodespace(
-	ctx context.Context, client apiClient, log *output.Logger, user *api.User, provisioningCodespace *api.Codespace,
-) (*api.Codespace, error) {
-	ticker := time.NewTicker(1 * time.Second)
+func pollForCodespace(ctx context.Context, client apiClient, log *output.Logger, duration, interval time.Duration, user, name string) (*api.Codespace, error) {
+	ctx, cancel := context.WithTimeout(ctx, duration)
+	defer cancel()
+
+	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
 	for {
@@ -140,18 +138,13 @@ func pollForCodespace(
 			return nil, ctx.Err()
 		case <-ticker.C:
 			log.Print(".")
-			token, err := client.GetCodespaceToken(ctx, user.Login, provisioningCodespace.Name)
+			token, err := client.GetCodespaceToken(ctx, user, name)
 			if err != nil {
 				// Do nothing. We expect this to fail until the codespace is provisioned
 				continue
 			}
 
-			codespace, err := client.GetCodespace(ctx, token, user.Login, provisioningCodespace.Name)
-			if err != nil {
-				return nil, fmt.Errorf("failed to get codespace: %w", err)
-			}
-
-			return codespace, nil
+			return client.GetCodespace(ctx, token, user, name)
 		}
 	}
 }
diff --git a/cmd/ghcs/create_test.go b/cmd/ghcs/create_test.go
index 3df900afc..e86fa00e6 100644
--- a/cmd/ghcs/create_test.go
+++ b/cmd/ghcs/create_test.go
@@ -37,18 +37,13 @@ func TestPollForCodespace(t *testing.T) {
 	user := &api.User{Login: "test"}
 	tmpCodespace := &api.Codespace{Name: "tmp-codespace"}
 	codespaceToken := "codespace-token"
+	ctx := context.Background()
 
-	ctxTimeout := 1 * time.Second
-	exceedTime := 2 * time.Second
-	exceedProvisioningTime := false
+	pollInterval := 50 * time.Millisecond
+	pollTimeout := 100 * time.Millisecond
 
 	api := &mockAPIClient{
 		getCodespaceToken: func(ctx context.Context, userLogin, codespace string) (string, error) {
-			if exceedProvisioningTime {
-				ticker := time.NewTicker(exceedTime)
-				defer ticker.Stop()
-				<-ticker.C
-			}
 			if userLogin != user.Login {
 				return "", fmt.Errorf("user does not match, got: %s, expected: %s", userLogin, user.Login)
 			}
@@ -71,10 +66,7 @@ func TestPollForCodespace(t *testing.T) {
 		},
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), ctxTimeout)
-	defer cancel()
-
-	codespace, err := pollForCodespace(ctx, api, logger, user, tmpCodespace)
+	codespace, err := pollForCodespace(ctx, api, logger, pollTimeout, pollInterval, user.Login, tmpCodespace.Name)
 	if err != nil {
 		t.Error(err)
 	}
@@ -82,12 +74,10 @@ func TestPollForCodespace(t *testing.T) {
 		t.Errorf("returned codespace does not match, got: %s, expected: %s", codespace.Name, tmpCodespace.Name)
 	}
 
-	exceedProvisioningTime = true
-	ctx, cancel = context.WithTimeout(ctx, ctxTimeout)
-	defer cancel()
-
-	_, err = pollForCodespace(ctx, api, logger, user, tmpCodespace)
-	if err == nil {
+	// swap the durations to trigger a timeout
+	pollTimeout, pollInterval = pollInterval, pollTimeout
+	_, err = pollForCodespace(ctx, api, logger, pollTimeout, pollInterval, user.Login, tmpCodespace.Name)
+	if err != context.DeadlineExceeded {
 		t.Error("expected context deadline exceeded error, got nil")
 	}
 }