chore: remove dbfake.WorkspaceWithAgent (#10879)

Replace dbfake.WorkspaceWithAgent() with the builder pattern and remove this function.

cli/agent_test.go

@@ -31,22 +31,25 @@ func TestWorkspaceAgent(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
-			OrganizationID: user.OrganizationID,
-			OwnerID:        user.UserID,
-		})
+		r := dbfake.NewWorkspaceBuilder(t, db).
+			Seed(database.Workspace{
+				OrganizationID: user.OrganizationID,
+				OwnerID:        user.UserID,
+			}).
+			WithAgent().
+			Do()
 		logDir := t.TempDir()
 		inv, _ := clitest.New(t,
 			"agent",
 			"--auth", "token",
-			"--agent-token", authToken,
+			"--agent-token", r.AgentToken,
 			"--agent-url", client.URL.String(),
 			"--log-dir", logDir,
 		)
 
 		clitest.Start(t, inv)
 
-		coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+		coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 
 		require.Eventually(t, func() bool {
 			info, err := os.Stat(filepath.Join(logDir, "coder-agent.log"))
@@ -209,16 +212,19 @@ func TestWorkspaceAgent(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
-			OrganizationID: user.OrganizationID,
-			OwnerID:        user.UserID,
-		})
+		r := dbfake.NewWorkspaceBuilder(t, db).
+			Seed(database.Workspace{
+				OrganizationID: user.OrganizationID,
+				OwnerID:        user.UserID,
+			}).
+			WithAgent().
+			Do()
 
 		logDir := t.TempDir()
 		inv, _ := clitest.New(t,
 			"agent",
 			"--auth", "token",
-			"--agent-token", authToken,
+			"--agent-token", r.AgentToken,
 			"--agent-url", client.URL.String(),
 			"--log-dir", logDir,
 		)
@@ -227,7 +233,7 @@ func TestWorkspaceAgent(t *testing.T) {
 
 		clitest.Start(t, inv)
 
-		resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+		resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 		require.Len(t, resources, 1)
 		require.Len(t, resources[0].Agents, 1)
 		require.Len(t, resources[0].Agents[0].Subsystems, 2)

cli/configssh_test.go

@@ -77,12 +77,15 @@ func TestConfigSSH(t *testing.T) {
 	})
 	owner := coderdtest.CreateFirstUser(t, client)
 	member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
-	ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
-		OrganizationID: owner.OrganizationID,
-		OwnerID:        memberUser.ID,
-	})
-	_ = agenttest.New(t, client.URL, authToken)
-	resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+	r := dbfake.NewWorkspaceBuilder(t, db).
+		Seed(database.Workspace{
+			OrganizationID: owner.OrganizationID,
+			OwnerID:        memberUser.ID,
+		}).
+		WithAgent().
+		Do()
+	_ = agenttest.New(t, client.URL, r.AgentToken)
+	resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 	agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
 	require.NoError(t, err)
 	defer agentConn.Close()
@@ -153,7 +156,7 @@ func TestConfigSSH(t *testing.T) {
 
 	home := filepath.Dir(filepath.Dir(sshConfigFile))
 	// #nosec
-	sshCmd := exec.Command("ssh", "-F", sshConfigFile, hostname+ws.Name, "echo", "test")
+	sshCmd := exec.Command("ssh", "-F", sshConfigFile, hostname+r.Workspace.Name, "echo", "test")
 	pty = ptytest.New(t)
 	// Set HOME because coder config is included from ~/.ssh/coder.
 	sshCmd.Env = append(sshCmd.Env, fmt.Sprintf("HOME=%s", home))
@@ -572,10 +575,10 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
 			client, db := coderdtest.NewWithDatabase(t, nil)
 			user := coderdtest.CreateFirstUser(t, client)
 			if tt.hasAgent {
-				_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+				_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 					OrganizationID: user.OrganizationID,
 					OwnerID:        user.UserID,
-				})
+				}).WithAgent().Do()
 			}
 
 			// Prepare ssh config files.

cli/gitssh_test.go

@@ -48,18 +48,21 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str
 	require.NoError(t, err)
 
 	// setup template
-	ws, agentToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
-		OrganizationID: user.OrganizationID,
-		OwnerID:        user.UserID,
-	})
+	r := dbfake.NewWorkspaceBuilder(t, db).
+		Seed(database.Workspace{
+			OrganizationID: user.OrganizationID,
+			OwnerID:        user.UserID,
+		}).
+		WithAgent().
+		Do()
 	// start workspace agent
 	agentClient := agentsdk.New(client.URL)
-	agentClient.SetSessionToken(agentToken)
-	_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
+	agentClient.SetSessionToken(r.AgentToken)
+	_ = agenttest.New(t, client.URL, r.AgentToken, func(o *agent.Options) {
 		o.Client = agentClient
 	})
-	_ = coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
-	return agentClient, agentToken, pubkey
+	_ = coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
+	return agentClient, r.AgentToken, pubkey
 }
 
 func serveSSHForGitSSH(t *testing.T, handler func(ssh.Session), pubkeys ...gossh.PublicKey) *net.TCPAddr {

cli/list_test.go

@@ -25,10 +25,10 @@ func TestList(t *testing.T) {
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		owner := coderdtest.CreateFirstUser(t, client)
 		member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
-		ws, _ := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: owner.OrganizationID,
 			OwnerID:        memberUser.ID,
-		})
+		}).WithAgent().Do()
 		inv, root := clitest.New(t, "ls")
 		clitest.SetupConfig(t, member, root)
 		pty := ptytest.New(t).Attach(inv)
@@ -41,7 +41,7 @@ func TestList(t *testing.T) {
 			assert.NoError(t, errC)
 			close(done)
 		}()
-		pty.ExpectMatch(ws.Name)
+		pty.ExpectMatch(r.Workspace.Name)
 		pty.ExpectMatch("Started")
 		cancelFunc()
 		<-done
@@ -52,10 +52,10 @@ func TestList(t *testing.T) {
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		owner := coderdtest.CreateFirstUser(t, client)
 		member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
-		dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: owner.OrganizationID,
 			OwnerID:        memberUser.ID,
-		})
+		}).WithAgent().Do()
 
 		inv, root := clitest.New(t, "list", "--output=json")
 		clitest.SetupConfig(t, member, root)

cli/portforward_test.go

@@ -305,17 +305,17 @@ func runAgent(t *testing.T, client *codersdk.Client, owner uuid.UUID, db databas
 	require.NoError(t, err, "specified user does not exist")
 	require.Greater(t, len(user.OrganizationIDs), 0, "user has no organizations")
 	orgID := user.OrganizationIDs[0]
-	ws, agentToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		OrganizationID: orgID,
 		OwnerID:        owner,
-	})
-	_ = agenttest.New(t, client.URL, agentToken,
+	}).WithAgent().Do()
+	_ = agenttest.New(t, client.URL, r.AgentToken,
 		func(o *agent.Options) {
 			o.SSHMaxTimeout = 60 * time.Second
 		},
 	)
-	coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
-	return ws
+	coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
+	return r.Workspace
 }
 
 // setupTestListener starts accepting connections and echoing a single packet.

cli/schedule_test.go

@@ -38,30 +38,30 @@ func setupTestSchedule(t *testing.T, sched *cron.Schedule) (ownerClient, memberC
 	memberClient, memberUser := coderdtest.CreateAnotherUserMutators(t, ownerClient, owner.OrganizationID, nil, func(r *codersdk.CreateUserRequest) {
 		r.Username = "testuser2" // ensure deterministic ordering
 	})
-	_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		Name:              "a-owner",
 		OwnerID:           owner.UserID,
 		OrganizationID:    owner.OrganizationID,
 		AutostartSchedule: sql.NullString{String: sched.String(), Valid: true},
 		Ttl:               sql.NullInt64{Int64: 8 * time.Hour.Nanoseconds(), Valid: true},
-	})
-	_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	}).WithAgent().Do()
+	_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		Name:              "b-owner",
 		OwnerID:           owner.UserID,
 		OrganizationID:    owner.OrganizationID,
 		AutostartSchedule: sql.NullString{String: sched.String(), Valid: true},
-	})
-	_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	}).WithAgent().Do()
+	_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		Name:           "c-member",
 		OwnerID:        memberUser.ID,
 		OrganizationID: owner.OrganizationID,
 		Ttl:            sql.NullInt64{Int64: 8 * time.Hour.Nanoseconds(), Valid: true},
-	})
-	_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	}).WithAgent().Do()
+	_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		Name:           "d-member",
 		OwnerID:        memberUser.ID,
 		OrganizationID: owner.OrganizationID,
-	})
+	}).WithAgent().Do()
 
 	// Need this for LatestBuild.Deadline
 	resp, err := ownerClient.Workspaces(context.Background(), codersdk.WorkspaceFilter{})

cli/ssh_test.go

@@ -52,12 +52,15 @@ func setupWorkspaceForAgent(t *testing.T, mutations ...func([]*proto.Agent) []*p
 	client.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
 	first := coderdtest.CreateFirstUser(t, client)
 	userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
-	workspace, agentToken := dbfake.WorkspaceWithAgent(t, store, database.Workspace{
-		OrganizationID: first.OrganizationID,
-		OwnerID:        user.ID,
-	}, mutations...)
+	r := dbfake.NewWorkspaceBuilder(t, store).
+		Seed(database.Workspace{
+			OrganizationID: first.OrganizationID,
+			OwnerID:        user.ID,
+		}).
+		WithAgent(mutations...).
+		Do()
 
-	return userClient, workspace, agentToken
+	return userClient, r.Workspace, r.AgentToken
 }
 
 func TestSSH(t *testing.T) {
@@ -127,11 +130,11 @@ func TestSSH(t *testing.T) {
 		client.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
 		first := coderdtest.CreateFirstUser(t, client)
 		userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
-		workspace, agentToken := dbfake.WorkspaceWithAgent(t, store, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, store).Seed(database.Workspace{
 			OrganizationID: first.OrganizationID,
 			OwnerID:        user.ID,
-		})
-		inv, root := clitest.New(t, "ssh", workspace.Name)
+		}).WithAgent().Do()
+		inv, root := clitest.New(t, "ssh", r.Workspace.Name)
 		clitest.SetupConfig(t, userClient, root)
 		pty := ptytest.New(t).Attach(inv)
 
@@ -144,14 +147,14 @@ func TestSSH(t *testing.T) {
 		})
 		pty.ExpectMatch("Waiting")
 
-		_ = agenttest.New(t, client.URL, agentToken)
-		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+		_ = agenttest.New(t, client.URL, r.AgentToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 
 		// Ensure the agent is connected.
 		pty.WriteLine("echo hell'o'")
 		pty.ExpectMatchContext(ctx, "hello")
 
-		_ = dbfake.NewWorkspaceBuildBuilder(t, store, workspace).
+		_ = dbfake.NewWorkspaceBuildBuilder(t, store, r.Workspace).
 			Seed(database.WorkspaceBuild{
 				Transition:  database.WorkspaceTransitionStop,
 				BuildNumber: 2,
@@ -466,15 +469,15 @@ func TestSSH(t *testing.T) {
 		client.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
 		first := coderdtest.CreateFirstUser(t, client)
 		userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
-		workspace, agentToken := dbfake.WorkspaceWithAgent(t, store, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, store).Seed(database.Workspace{
 			OrganizationID: first.OrganizationID,
 			OwnerID:        user.ID,
-		})
+		}).WithAgent().Do()
 
 		_, _ = tGoContext(t, func(ctx context.Context) {
 			// Run this async so the SSH command has to wait for
 			// the build and agent to connect.
-			_ = agenttest.New(t, client.URL, agentToken)
+			_ = agenttest.New(t, client.URL, r.AgentToken)
 			<-ctx.Done()
 		})
 
@@ -489,7 +492,7 @@ func TestSSH(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 
-		inv, root := clitest.New(t, "ssh", "--stdio", workspace.Name)
+		inv, root := clitest.New(t, "ssh", "--stdio", r.Workspace.Name)
 		clitest.SetupConfig(t, userClient, root)
 		inv.Stdin = clientOutput
 		inv.Stdout = serverInput
@@ -520,7 +523,7 @@ func TestSSH(t *testing.T) {
 		err = session.Shell()
 		require.NoError(t, err)
 
-		_ = dbfake.NewWorkspaceBuildBuilder(t, store, workspace).
+		_ = dbfake.NewWorkspaceBuildBuilder(t, store, r.Workspace).
 			Seed(database.WorkspaceBuild{
 				Transition:  database.WorkspaceTransitionStop,
 				BuildNumber: 2,

coderd/database/dbfake/dbfake.go

@@ -96,19 +96,6 @@ func Workspace(t testing.TB, db database.Store, seed database.Workspace) databas
 	return r.Workspace
 }
 
-// WorkspaceWithAgent is a helper that generates a workspace with a single resource
-// that has an agent attached to it. The agent token is returned.
-func WorkspaceWithAgent(
-	t testing.TB, db database.Store, seed database.Workspace,
-	mutations ...func([]*sdkproto.Agent) []*sdkproto.Agent,
-) (
-	database.Workspace, string,
-) {
-	t.Helper()
-	r := NewWorkspaceBuilder(t, db).Seed(seed).WithAgent(mutations...).Do()
-	return r.Workspace, r.AgentToken
-}
-
 type WorkspaceBuildBuilder struct {
 	t         testing.TB
 	db        database.Store

coderd/workspaceagents_test.go

@@ -224,13 +224,13 @@ func TestWorkspaceAgentLogs(t *testing.T) {
 		ctx := testutil.Context(t, testutil.WaitMedium)
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
+		}).WithAgent().Do()
 
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 		err := agentClient.PatchLogs(ctx, agentsdk.PatchLogs{
 			Logs: []agentsdk.Log{
 				{
@@ -244,7 +244,7 @@ func TestWorkspaceAgentLogs(t *testing.T) {
 			},
 		})
 		require.NoError(t, err)
-		workspace, err := client.Workspace(ctx, ws.ID)
+		workspace, err := client.Workspace(ctx, r.Workspace.ID)
 		require.NoError(t, err)
 		logs, closer, err := client.WorkspaceAgentLogsAfter(ctx, workspace.LatestBuild.Resources[0].Agents[0].ID, 0, true)
 		require.NoError(t, err)
@@ -266,12 +266,12 @@ func TestWorkspaceAgentLogs(t *testing.T) {
 		ctx := testutil.Context(t, testutil.WaitMedium)
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
+		}).WithAgent().Do()
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 		err := agentClient.PatchLogs(ctx, agentsdk.PatchLogs{
 			Logs: []agentsdk.Log{
 				{
@@ -281,7 +281,7 @@ func TestWorkspaceAgentLogs(t *testing.T) {
 			},
 		})
 		require.NoError(t, err)
-		workspace, err := client.Workspace(ctx, ws.ID)
+		workspace, err := client.Workspace(ctx, r.Workspace.ID)
 		require.NoError(t, err)
 		logs, closer, err := client.WorkspaceAgentLogsAfter(ctx, workspace.LatestBuild.Resources[0].Agents[0].ID, 0, true)
 		require.NoError(t, err)
@@ -308,15 +308,15 @@ func TestWorkspaceAgentLogs(t *testing.T) {
 		ctx := testutil.Context(t, testutil.WaitMedium)
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
-		updates, err := client.WatchWorkspace(ctx, ws.ID)
+		}).WithAgent().Do()
+		updates, err := client.WatchWorkspace(ctx, r.Workspace.ID)
 		require.NoError(t, err)
 
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 		err = agentClient.PatchLogs(ctx, agentsdk.PatchLogs{
 			Logs: []agentsdk.Log{{
 				CreatedAt: dbtime.Now(),
@@ -351,12 +351,12 @@ func TestWorkspaceAgentListen(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
-		_ = agenttest.New(t, client.URL, authToken)
-		resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+		}).WithAgent().Do()
+		_ = agenttest.New(t, client.URL, r.AgentToken)
+		resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
@@ -437,13 +437,13 @@ func TestWorkspaceAgentTailnet(t *testing.T) {
 	client, db := coderdtest.NewWithDatabase(t, nil)
 	user := coderdtest.CreateFirstUser(t, client)
 
-	ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		OrganizationID: user.OrganizationID,
 		OwnerID:        user.UserID,
-	})
+	}).WithAgent().Do()
 
-	_ = agenttest.New(t, client.URL, authToken)
-	resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+	_ = agenttest.New(t, client.URL, r.AgentToken)
+	resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	defer cancelFunc()
@@ -476,21 +476,21 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) {
 		DeploymentValues: dv,
 	})
 	user := coderdtest.CreateFirstUser(t, client)
-	ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+	r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 		OrganizationID: user.OrganizationID,
 		OwnerID:        user.UserID,
-	})
+	}).WithAgent().Do()
 	ctx := testutil.Context(t, testutil.WaitLong)
 
 	// Verify that the manifest has DisableDirectConnections set to true.
 	agentClient := agentsdk.New(client.URL)
-	agentClient.SetSessionToken(authToken)
+	agentClient.SetSessionToken(r.AgentToken)
 	manifest, err := agentClient.Manifest(ctx)
 	require.NoError(t, err)
 	require.True(t, manifest.DisableDirectConnections)
 
-	_ = agenttest.New(t, client.URL, authToken)
-	resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+	_ = agenttest.New(t, client.URL, r.AgentToken)
+	resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 	agentID := resources[0].Agents[0].ID
 
 	// Verify that the connection data has no STUN ports and
@@ -856,13 +856,13 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
+		}).WithAgent().Do()
 
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 
 		_, err := agentClient.PostStats(context.Background(), &agentsdk.Stats{
 			ConnectionsByProto: map[string]int64{"TCP": 1},
@@ -882,12 +882,12 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		newWorkspace, err := client.Workspace(context.Background(), ws.ID)
+		newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
 		require.NoError(t, err)
 
 		assert.True(t,
-			newWorkspace.LastUsedAt.Equal(ws.LastUsedAt),
-			"%s and %s should not differ", newWorkspace.LastUsedAt, ws.LastUsedAt,
+			newWorkspace.LastUsedAt.Equal(r.Workspace.LastUsedAt),
+			"%s and %s should not differ", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
 		)
 
 		_, err = agentClient.PostStats(context.Background(), &agentsdk.Stats{
@@ -905,12 +905,12 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		newWorkspace, err = client.Workspace(context.Background(), ws.ID)
+		newWorkspace, err = client.Workspace(context.Background(), r.Workspace.ID)
 		require.NoError(t, err)
 
 		assert.True(t,
-			newWorkspace.LastUsedAt.After(ws.LastUsedAt),
-			"%s is not after %s", newWorkspace.LastUsedAt, ws.LastUsedAt,
+			newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
+			"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
 		)
 	})
 }
@@ -923,11 +923,11 @@ func TestWorkspaceAgent_LifecycleState(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
-		workspace, err := client.Workspace(context.Background(), ws.ID)
+		}).WithAgent().Do()
+		workspace, err := client.Workspace(context.Background(), r.Workspace.ID)
 		require.NoError(t, err)
 		for _, res := range workspace.LatestBuild.Resources {
 			for _, a := range res.Agents {
@@ -936,7 +936,7 @@ func TestWorkspaceAgent_LifecycleState(t *testing.T) {
 		}
 
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 
 		tests := []struct {
 			state   codersdk.WorkspaceAgentLifecycle
@@ -1312,12 +1312,12 @@ func TestWorkspaceAgent_Startup(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
+		}).WithAgent().Do()
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 
 		ctx := testutil.Context(t, testutil.WaitMedium)
 
@@ -1341,7 +1341,7 @@ func TestWorkspaceAgent_Startup(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		workspace, err := client.Workspace(ctx, ws.ID)
+		workspace, err := client.Workspace(ctx, r.Workspace.ID)
 		require.NoError(t, err)
 
 		wsagent, err := client.WorkspaceAgent(ctx, workspace.LatestBuild.Resources[0].Agents[0].ID)
@@ -1358,13 +1358,13 @@ func TestWorkspaceAgent_Startup(t *testing.T) {
 
 		client, db := coderdtest.NewWithDatabase(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		_, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
+		r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
 			OrganizationID: user.OrganizationID,
 			OwnerID:        user.UserID,
-		})
+		}).WithAgent().Do()
 
 		agentClient := agentsdk.New(client.URL)
-		agentClient.SetSessionToken(authToken)
+		agentClient.SetSessionToken(r.AgentToken)
 
 		ctx := testutil.Context(t, testutil.WaitMedium)
 
@@ -1406,13 +1406,13 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) {
 	api.DERPMapper.Store(&derpMapFn)
 
 	// Start workspace a workspace agent.
-	ws, agentToken := dbfake.WorkspaceWithAgent(t, api.Database, database.Workspace{
+	r := dbfake.NewWorkspaceBuilder(t, api.Database).Seed(database.Workspace{
 		OrganizationID: user.OrganizationID,
 		OwnerID:        user.UserID,
-	})
+	}).WithAgent().Do()
 
-	agentCloser := agenttest.New(t, client.URL, agentToken)
-	resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
+	agentCloser := agenttest.New(t, client.URL, r.AgentToken)
+	resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
 	agentID := resources[0].Agents[0].ID
 
 	// Connect from a client.