From b90c74a94d0c271a5d73ace01281941dac2acd4f Mon Sep 17 00:00:00 2001
From: Ethan <39577870+ethanndickson@users.noreply.github.com>
Date: Thu, 30 Oct 2025 12:14:25 +1100
Subject: [PATCH] chore(scaletest): avoid polling workspace builds during
 workspace-updates tests (#20534)

This PR is just committing the changes I made while running the
`workspace-updates` load generator.

It ensures we're not polling the workspace build progress in the
background (while we also watch for workspace updates via the tailnet),
and also removes an unnecessary query to `/api/v2/workspace/{id}` after
each workspace is built.
---
 scaletest/createworkspaces/run.go |  6 +++++-
 scaletest/workspacebuild/run.go   | 26 ++++++++++++++------------
 scaletest/workspaceupdates/run.go |  4 ++++
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/scaletest/createworkspaces/run.go b/scaletest/createworkspaces/run.go
index 49fe0548b3..09903c06cf 100644
--- a/scaletest/createworkspaces/run.go
+++ b/scaletest/createworkspaces/run.go
@@ -87,10 +87,14 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 	workspaceBuildConfig.OrganizationID = r.cfg.User.OrganizationID
 	workspaceBuildConfig.UserID = user.ID.String()
 	r.workspacebuildRunner = workspacebuild.NewRunner(client, workspaceBuildConfig)
-	workspace, err := r.workspacebuildRunner.RunReturningWorkspace(ctx, id, logs)
+	slimWorkspace, err := r.workspacebuildRunner.RunReturningWorkspace(ctx, id, logs)
 	if err != nil {
 		return xerrors.Errorf("create workspace: %w", err)
 	}
+	workspace, err := client.Workspace(ctx, slimWorkspace.ID)
+	if err != nil {
+		return xerrors.Errorf("get full workspace info: %w", err)
+	}
 
 	if r.cfg.Workspace.NoWaitForAgents {
 		return nil
diff --git a/scaletest/workspacebuild/run.go b/scaletest/workspacebuild/run.go
index 308c18f0b6..f05173c688 100644
--- a/scaletest/workspacebuild/run.go
+++ b/scaletest/workspacebuild/run.go
@@ -33,8 +33,13 @@ func NewRunner(client *codersdk.Client, cfg Config) *Runner {
 	}
 }
 
+type SlimWorkspace struct {
+	ID   uuid.UUID
+	Name string
+}
+
 // Run implements Runnable.
-func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.Writer) (codersdk.Workspace, error) {
+func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.Writer) (SlimWorkspace, error) {
 	ctx, span := tracing.StartSpan(ctx)
 	defer span.End()
 
@@ -47,14 +52,14 @@ func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.W
 	if req.Name == "" {
 		randName, err := loadtestutil.GenerateWorkspaceName(id)
 		if err != nil {
-			return codersdk.Workspace{}, xerrors.Errorf("generate random name for workspace: %w", err)
+			return SlimWorkspace{}, xerrors.Errorf("generate random name for workspace: %w", err)
 		}
 		req.Name = randName
 	}
 
 	workspace, err := r.client.CreateWorkspace(ctx, r.cfg.OrganizationID, r.cfg.UserID, req)
 	if err != nil {
-		return codersdk.Workspace{}, xerrors.Errorf("create workspace: %w", err)
+		return SlimWorkspace{}, xerrors.Errorf("create workspace: %w", err)
 	}
 	r.workspaceID = workspace.ID
 
@@ -72,7 +77,7 @@ func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.W
 					TemplateVersionID:   req.TemplateVersionID,
 				})
 				if err != nil {
-					return codersdk.Workspace{}, xerrors.Errorf("create workspace build: %w", err)
+					return SlimWorkspace{}, xerrors.Errorf("create workspace build: %w", err)
 				}
 				err = waitForBuild(ctx, logs, r.client, workspace.LatestBuild.ID)
 				if err == nil {
@@ -80,7 +85,7 @@ func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.W
 				}
 			}
 			if err != nil {
-				return codersdk.Workspace{}, xerrors.Errorf("wait for build: %w", err)
+				return SlimWorkspace{}, xerrors.Errorf("wait for build: %w", err)
 			}
 		}
 	}
@@ -91,16 +96,13 @@ func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.W
 		_, _ = fmt.Fprintln(logs, "")
 		err = waitForAgents(ctx, logs, r.client, workspace.ID)
 		if err != nil {
-			return codersdk.Workspace{}, xerrors.Errorf("wait for agent: %w", err)
+			return SlimWorkspace{}, xerrors.Errorf("wait for agent: %w", err)
 		}
 	}
 
-	workspace, err = r.client.Workspace(ctx, workspace.ID)
-	if err != nil {
-		return codersdk.Workspace{}, xerrors.Errorf("get workspace %q: %w", workspace.ID.String(), err)
-	}
-
-	return workspace, nil
+	// Some users of this runner might not need the full workspace, and
+	// want to avoid querying the workspace.
+	return SlimWorkspace{ID: workspace.ID, Name: workspace.Name}, nil
 }
 
 // CleanupRunner is a runner that deletes a workspace in the Run phase.
diff --git a/scaletest/workspaceupdates/run.go b/scaletest/workspaceupdates/run.go
index 4addf2b5a5..fa05d290f0 100644
--- a/scaletest/workspaceupdates/run.go
+++ b/scaletest/workspaceupdates/run.go
@@ -116,6 +116,10 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 		workspaceBuildConfig.OrganizationID = r.cfg.User.OrganizationID
 		workspaceBuildConfig.UserID = newUser.ID.String()
 		workspaceBuildConfig.Request.Name = workspaceName
+		// We'll watch for completion ourselves via the tailnet workspace
+		// updates stream.
+		workspaceBuildConfig.NoWaitForAgents = true
+		workspaceBuildConfig.NoWaitForBuild = true
 
 		runner := workspacebuild.NewRunner(newUserClient, workspaceBuildConfig)
 		r.workspacebuildRunners = append(r.workspacebuildRunners, runner)