From cb6552f4cae0d5255471319455f24254a5bcfae0 Mon Sep 17 00:00:00 2001
From: Jose Garcia <jose.balius@gmail.com>
Date: Fri, 27 Aug 2021 12:50:32 -0400
Subject: [PATCH] more efficient impl for processing states

---
 cmd/ghcs/create.go                | 39 ++++++++++++++++++-------------
 internal/codespaces/codespaces.go | 11 ++++++++-
 internal/codespaces/states.go     |  5 ++++
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/cmd/ghcs/create.go b/cmd/ghcs/create.go
index fb9834351..9ffcc58d1 100644
--- a/cmd/ghcs/create.go
+++ b/cmd/ghcs/create.go
@@ -112,6 +112,7 @@ func showStatus(ctx context.Context, log *output.Logger, apiClient *api.API, use
 	}
 
 	var lastState codespaces.PostCreateState
+	finishedStates := make(map[string]bool)
 	var breakNextState bool
 
 	for {
@@ -122,25 +123,31 @@ func showStatus(ctx context.Context, log *output.Logger, apiClient *api.API, use
 
 		var inProgress bool
 		for _, state := range stateUpdate.PostCreateStates {
-			switch state.Status {
-			case codespaces.PostCreateStateRunning:
-				if lastState != state {
-					lastState = state
-					log.Print(state.Name)
-				} else {
-					log.Print(".")
-				}
+			if _, found := finishedStates[state.Name]; found {
+				continue // skip this state as we've processed it already
+			}
 
-				inProgress = true
-			case codespaces.PostCreateStateFailed:
-				if lastState.Name == state.Name && lastState.Status != state.Status {
+			if state.Name != lastState.Name {
+				log.Print(state.Name)
+
+				if state.Status == codespaces.PostCreateStateRunning {
+					inProgress = true
 					lastState = state
-					log.Print(".Failed\n")
+					log.Print("...")
+					break
+				} else {
+					finishedStates[state.Name] = true
+					log.Println("..." + state.Status)
 				}
-			case codespaces.PostCreateStateSuccess:
-				if lastState.Name == state.Name && lastState.Status != state.Status {
-					lastState = state
-					log.Print(".Success\n")
+			} else {
+				if state.Status == codespaces.PostCreateStateRunning {
+					inProgress = true
+					log.Print(".")
+					break
+				} else {
+					finishedStates[state.Name] = true
+					log.Println(state.Status)
+					lastState = codespaces.PostCreateState{} // reset the value
 				}
 			}
 		}
diff --git a/internal/codespaces/codespaces.go b/internal/codespaces/codespaces.go
index 005ea0fda..11e0a8902 100644
--- a/internal/codespaces/codespaces.go
+++ b/internal/codespaces/codespaces.go
@@ -62,6 +62,15 @@ type logger interface {
 	Println(v ...interface{}) (int, error)
 }
 
+func connectionReady(codespace *api.Codespace) bool {
+	ready := codespace.Environment.Connection.SessionID != ""
+	ready = ready && codespace.Environment.Connection.SessionToken != ""
+	ready = ready && codespace.Environment.Connection.RelayEndpoint != ""
+	ready = ready && codespace.Environment.Connection.RelaySAS != ""
+	ready = ready && codespace.Environment.State == api.CodespaceEnvironmentStateAvailable
+	return ready
+}
+
 func ConnectToLiveshare(ctx context.Context, log logger, apiClient *api.API, userLogin, token string, codespace *api.Codespace) (client *liveshare.Client, err error) {
 	var startedCodespace bool
 	if codespace.Environment.State != api.CodespaceEnvironmentStateAvailable {
@@ -73,7 +82,7 @@ func ConnectToLiveshare(ctx context.Context, log logger, apiClient *api.API, use
 	}
 
 	retries := 0
-	for codespace.Environment.Connection.SessionID == "" || codespace.Environment.State != api.CodespaceEnvironmentStateAvailable {
+	for !connectionReady(codespace) {
 		if retries > 1 {
 			if retries%2 == 0 {
 				log.Print(".")
diff --git a/internal/codespaces/states.go b/internal/codespaces/states.go
index 078dc546c..3ad6c9a4c 100644
--- a/internal/codespaces/states.go
+++ b/internal/codespaces/states.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"strings"
 	"time"
 
 	"github.com/github/ghcs/api"
@@ -12,6 +13,10 @@ import (
 
 type PostCreateStateStatus string
 
+func (p PostCreateStateStatus) String() string {
+	return strings.Title(string(p))
+}
+
 const (
 	PostCreateStateRunning PostCreateStateStatus = "running"
 	PostCreateStateSuccess PostCreateStateStatus = "succeeded"