refactor(scaletest): make workspacebuild return the built workspace (#19997)

Similar idea as in #19811, this runner doesn't need to conform to `Runnable`, so we have it return the workspace from the `RunReturningWorkspace` function, instead of the more fragile `Run`, followed by a `.WorkspaceID()`.

scaletest/createworkspaces/run.go

@@ -87,7 +87,7 @@ 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)
-	err = r.workspacebuildRunner.Run(ctx, id, logs)
+	workspace, err := r.workspacebuildRunner.RunReturningWorkspace(ctx, id, logs)
 	if err != nil {
 		return xerrors.Errorf("create workspace: %w", err)
 	}
@@ -96,16 +96,6 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 		return nil
 	}
 
-	// Get the workspace.
-	workspaceID, err := r.workspacebuildRunner.WorkspaceID()
-	if err != nil {
-		return xerrors.Errorf("get workspace ID: %w", err)
-	}
-	workspace, err := client.Workspace(ctx, workspaceID)
-	if err != nil {
-		return xerrors.Errorf("get workspace %q: %w", workspaceID.String(), err)
-	}
-
 	// Find the first agent.
 	var agent codersdk.WorkspaceAgent
 resourceLoop:
@@ -116,7 +106,7 @@ resourceLoop:
 		}
 	}
 	if agent.ID == uuid.Nil {
-		return xerrors.Errorf("no agents found for workspace %q", workspaceID.String())
+		return xerrors.Errorf("no agents found for workspace %q", workspace.ID.String())
 	}
 
 	eg, egCtx := errgroup.WithContext(ctx)

scaletest/workspacebuild/run.go

@@ -26,11 +26,6 @@ type Runner struct {
 	workspaceID uuid.UUID
 }
 
-var (
-	_ harness.Runnable  = &Runner{}
-	_ harness.Cleanable = &Runner{}
-)
-
 func NewRunner(client *codersdk.Client, cfg Config) *Runner {
 	return &Runner{
 		client: client,
@@ -39,7 +34,7 @@ func NewRunner(client *codersdk.Client, cfg Config) *Runner {
 }
 
 // Run implements Runnable.
-func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
+func (r *Runner) RunReturningWorkspace(ctx context.Context, id string, logs io.Writer) (codersdk.Workspace, error) {
 	ctx, span := tracing.StartSpan(ctx)
 	defer span.End()
 
@@ -52,14 +47,14 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 	if req.Name == "" {
 		randName, err := loadtestutil.GenerateWorkspaceName(id)
 		if err != nil {
-			return xerrors.Errorf("generate random name for workspace: %w", err)
+			return codersdk.Workspace{}, 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 xerrors.Errorf("create workspace: %w", err)
+		return codersdk.Workspace{}, xerrors.Errorf("create workspace: %w", err)
 	}
 	r.workspaceID = workspace.ID
 
@@ -74,7 +69,7 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 				TemplateVersionID:   req.TemplateVersionID,
 			})
 			if err != nil {
-				return xerrors.Errorf("create workspace build: %w", err)
+				return codersdk.Workspace{}, xerrors.Errorf("create workspace build: %w", err)
 			}
 			err = waitForBuild(ctx, logs, r.client, workspace.LatestBuild.ID)
 			if err == nil {
@@ -82,7 +77,7 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 			}
 		}
 		if err != nil {
-			return xerrors.Errorf("wait for build: %w", err)
+			return codersdk.Workspace{}, xerrors.Errorf("wait for build: %w", err)
 		}
 	}
 
@@ -92,19 +87,16 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 		_, _ = fmt.Fprintln(logs, "")
 		err = waitForAgents(ctx, logs, r.client, workspace.ID)
 		if err != nil {
-			return xerrors.Errorf("wait for agent: %w", err)
+			return codersdk.Workspace{}, xerrors.Errorf("wait for agent: %w", err)
 		}
 	}
 
-	return nil
-}
-
-func (r *Runner) WorkspaceID() (uuid.UUID, error) {
-	if r.workspaceID == uuid.Nil {
-		return uuid.Nil, xerrors.New("workspace ID not set")
+	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 r.workspaceID, nil
+	return workspace, nil
 }
 
 // CleanupRunner is a runner that deletes a workspace in the Run phase.

scaletest/workspacebuild/run_test.go

@@ -158,7 +158,7 @@ func Test_Runner(t *testing.T) {
 		})
 
 		logs := bytes.NewBuffer(nil)
-		err := runner.Run(ctx, "1", logs)
+		_, err := runner.RunReturningWorkspace(ctx, "1", logs)
 		logsStr := logs.String()
 		t.Log("Runner logs:\n\n" + logsStr)
 		require.NoError(t, err)
@@ -224,7 +224,7 @@ func Test_Runner(t *testing.T) {
 		})
 
 		logs := bytes.NewBuffer(nil)
-		err := runner.Run(ctx, "1", logs)
+		_, err := runner.RunReturningWorkspace(ctx, "1", logs)
 		logsStr := logs.String()
 		t.Log("Runner logs:\n\n" + logsStr)
 		require.Error(t, err)
@@ -271,7 +271,7 @@ func Test_Runner(t *testing.T) {
 		})
 
 		logs := bytes.NewBuffer(nil)
-		err := runner.Run(ctx, "1", logs)
+		_, err := runner.RunReturningWorkspace(ctx, "1", logs)
 		logsStr := logs.String()
 		t.Log("Runner logs:\n\n" + logsStr)
 		require.Error(t, err)

scaletest/workspaceupdates/run.go

@@ -126,7 +126,7 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error {
 		r.workspaces[workspaceName] = &workspace{
 			buildStartTime: time.Now(),
 		}
-		err = runner.Run(ctx, fmt.Sprintf("%s-%d", id, i), logs)
+		_, err = runner.RunReturningWorkspace(ctx, fmt.Sprintf("%s-%d", id, i), logs)
 		if err != nil {
 			return xerrors.Errorf("create workspace %d: %w", i, err)
 		}