improve synchro

mafredri · mafredri · commit 3954f5bf143e · 2025-01-09T09:49:26.000Z
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -155,17 +155,19 @@ func TestSSH(t *testing.T) {
 		isFirstBuild := true
 		buildURL := regexp.MustCompile("/api/v2/workspaces/.*/builds")
 		buildSync := make(chan struct{})
+		buildDone := make(chan struct{})
 		buildResume := make(chan struct{})
 		buildSyncMW := func(next http.Handler) http.Handler {
 			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				if r.Method == http.MethodPost && buildURL.MatchString(r.URL.Path) {
 					if !isFirstBuild {
+						t.Log("buildSyncMW: blocking post-build")
+						<-buildSync
 						defer func() {
-							t.Log("buildSyncMW: blocking post-build")
-							buildSync <- struct{}{}
+							buildDone <- struct{}{}
 							<-buildResume
-							t.Log("buildSyncMW: resuming...")
 						}()
+						t.Log("buildSyncMW: resuming...")
 					} else {
 						isFirstBuild = false
 					}
@@ -213,9 +215,19 @@ func TestSSH(t *testing.T) {
 		for _, pty := range ptys {
 			pty.ExpectMatchContext(ctx, "Workspace was stopped, starting workspace to allow connecting to")
 		}
-		for range ptys {
-			testutil.RequireRecvCtx(ctx, t, buildSync)
+
+		// Allow one build to complete.
+		testutil.RequireSendCtx(ctx, t, buildSync, struct{}{})
+		testutil.RequireRecvCtx(ctx, t, buildDone)
+
+		// Allow the remaining builds to continue.
+		for i := 0; i < len(ptys)-1; i++ {
+			testutil.RequireSendCtx(ctx, t, buildSync, struct{}{})
+		}
+		for i := 0; i < len(ptys)-1; i++ {
+			testutil.RequireRecvCtx(ctx, t, buildDone)
 		}
+		// Allow all three endpoints to return.
 		close(buildResume)
 
 		var foundConflict int
