Skip to content

Commit 23ad063

Browse files
spikecurtisspikecurtis
committed
chore: remove dbfake.WorkspaceWithAgent
1 parent 19c68a0 commit 23ad063

9 files changed

+128
-125
lines changed

‎cli/agent_test.go

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -31,22 +31,25 @@ func TestWorkspaceAgent(t *testing.T) {
3131

3232
client, db := coderdtest.NewWithDatabase(t, nil)
3333
user := coderdtest.CreateFirstUser(t, client)
34-
ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
35-
OrganizationID: user.OrganizationID,
36-
OwnerID: user.UserID,
37-
})
34+
r := dbfake.NewWorkspaceBuilder(t, db).
35+
Seed(database.Workspace{
36+
OrganizationID: user.OrganizationID,
37+
OwnerID: user.UserID,
38+
}).
39+
WithAgent().
40+
Do()
3841
logDir := t.TempDir()
3942
inv, _ := clitest.New(t,
4043
"agent",
4144
"--auth", "token",
42-
"--agent-token", authToken,
45+
"--agent-token", r.AgentToken,
4346
"--agent-url", client.URL.String(),
4447
"--log-dir", logDir,
4548
)
4649

4750
clitest.Start(t, inv)
4851

49-
coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
52+
coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
5053

5154
require.Eventually(t, func() bool {
5255
info, err := os.Stat(filepath.Join(logDir, "coder-agent.log"))
@@ -209,16 +212,19 @@ func TestWorkspaceAgent(t *testing.T) {
209212

210213
client, db := coderdtest.NewWithDatabase(t, nil)
211214
user := coderdtest.CreateFirstUser(t, client)
212-
ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
213-
OrganizationID: user.OrganizationID,
214-
OwnerID: user.UserID,
215-
})
215+
r := dbfake.NewWorkspaceBuilder(t, db).
216+
Seed(database.Workspace{
217+
OrganizationID: user.OrganizationID,
218+
OwnerID: user.UserID,
219+
}).
220+
WithAgent().
221+
Do()
216222

217223
logDir := t.TempDir()
218224
inv, _ := clitest.New(t,
219225
"agent",
220226
"--auth", "token",
221-
"--agent-token", authToken,
227+
"--agent-token", r.AgentToken,
222228
"--agent-url", client.URL.String(),
223229
"--log-dir", logDir,
224230
)
@@ -227,7 +233,7 @@ func TestWorkspaceAgent(t *testing.T) {
227233

228234
clitest.Start(t, inv)
229235

230-
resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
236+
resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
231237
require.Len(t, resources, 1)
232238
require.Len(t, resources[0].Agents, 1)
233239
require.Len(t, resources[0].Agents[0].Subsystems, 2)

‎cli/configssh_test.go

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -77,12 +77,15 @@ func TestConfigSSH(t *testing.T) {
7777
})
7878
owner := coderdtest.CreateFirstUser(t, client)
7979
member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
80-
ws, authToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
81-
OrganizationID: owner.OrganizationID,
82-
OwnerID: memberUser.ID,
83-
})
84-
_ = agenttest.New(t, client.URL, authToken)
85-
resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
80+
r := dbfake.NewWorkspaceBuilder(t, db).
81+
Seed(database.Workspace{
82+
OrganizationID: owner.OrganizationID,
83+
OwnerID: memberUser.ID,
84+
}).
85+
WithAgent().
86+
Do()
87+
_ = agenttest.New(t, client.URL, r.AgentToken)
88+
resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
8689
agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
8790
require.NoError(t, err)
8891
defer agentConn.Close()
@@ -153,7 +156,7 @@ func TestConfigSSH(t *testing.T) {
153156

154157
home := filepath.Dir(filepath.Dir(sshConfigFile))
155158
// #nosec
156-
sshCmd := exec.Command("ssh", "-F", sshConfigFile, hostname+ws.Name, "echo", "test")
159+
sshCmd := exec.Command("ssh", "-F", sshConfigFile, hostname+r.Workspace.Name, "echo", "test")
157160
pty = ptytest.New(t)
158161
// Set HOME because coder config is included from ~/.ssh/coder.
159162
sshCmd.Env = append(sshCmd.Env, fmt.Sprintf("HOME=%s", home))
@@ -572,10 +575,10 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
572575
client, db := coderdtest.NewWithDatabase(t, nil)
573576
user := coderdtest.CreateFirstUser(t, client)
574577
if tt.hasAgent {
575-
_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
578+
_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
576579
OrganizationID: user.OrganizationID,
577580
OwnerID: user.UserID,
578-
})
581+
}).WithAgent().Do()
579582
}
580583

581584
// Prepare ssh config files.

‎cli/gitssh_test.go

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -48,18 +48,21 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str
4848
require.NoError(t, err)
4949

5050
// setup template
51-
ws, agentToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
52-
OrganizationID: user.OrganizationID,
53-
OwnerID: user.UserID,
54-
})
51+
r := dbfake.NewWorkspaceBuilder(t, db).
52+
Seed(database.Workspace{
53+
OrganizationID: user.OrganizationID,
54+
OwnerID: user.UserID,
55+
}).
56+
WithAgent().
57+
Do()
5558
// start workspace agent
5659
agentClient := agentsdk.New(client.URL)
57-
agentClient.SetSessionToken(agentToken)
58-
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
60+
agentClient.SetSessionToken(r.AgentToken)
61+
_ = agenttest.New(t, client.URL, r.AgentToken, func(o *agent.Options) {
5962
o.Client = agentClient
6063
})
61-
_ = coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
62-
return agentClient, agentToken, pubkey
64+
_ = coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
65+
return agentClient, r.AgentToken, pubkey
6366
}
6467

6568
func serveSSHForGitSSH(t *testing.T, handler func(ssh.Session), pubkeys ...gossh.PublicKey) *net.TCPAddr {

‎cli/list_test.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,10 @@ func TestList(t *testing.T) {
2525
client, db := coderdtest.NewWithDatabase(t, nil)
2626
owner := coderdtest.CreateFirstUser(t, client)
2727
member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
28-
ws, _ := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
28+
r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
2929
OrganizationID: owner.OrganizationID,
3030
OwnerID: memberUser.ID,
31-
})
31+
}).WithAgent().Do()
3232
inv, root := clitest.New(t, "ls")
3333
clitest.SetupConfig(t, member, root)
3434
pty := ptytest.New(t).Attach(inv)
@@ -41,7 +41,7 @@ func TestList(t *testing.T) {
4141
assert.NoError(t, errC)
4242
close(done)
4343
}()
44-
pty.ExpectMatch(ws.Name)
44+
pty.ExpectMatch(r.Workspace.Name)
4545
pty.ExpectMatch("Started")
4646
cancelFunc()
4747
<-done
@@ -52,10 +52,10 @@ func TestList(t *testing.T) {
5252
client, db := coderdtest.NewWithDatabase(t, nil)
5353
owner := coderdtest.CreateFirstUser(t, client)
5454
member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
55-
dbfake.WorkspaceWithAgent(t, db, database.Workspace{
55+
dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
5656
OrganizationID: owner.OrganizationID,
5757
OwnerID: memberUser.ID,
58-
})
58+
}).WithAgent().Do()
5959

6060
inv, root := clitest.New(t, "list", "--output=json")
6161
clitest.SetupConfig(t, member, root)

‎cli/portforward_test.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -305,17 +305,17 @@ func runAgent(t *testing.T, client *codersdk.Client, owner uuid.UUID, db databas
305305
require.NoError(t, err, "specified user does not exist")
306306
require.Greater(t, len(user.OrganizationIDs), 0, "user has no organizations")
307307
orgID := user.OrganizationIDs[0]
308-
ws, agentToken := dbfake.WorkspaceWithAgent(t, db, database.Workspace{
308+
r := dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
309309
OrganizationID: orgID,
310310
OwnerID: owner,
311-
})
312-
_ = agenttest.New(t, client.URL, agentToken,
311+
}).WithAgent().Do()
312+
_ = agenttest.New(t, client.URL, r.AgentToken,
313313
func(o *agent.Options) {
314314
o.SSHMaxTimeout = 60 * time.Second
315315
},
316316
)
317-
coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
318-
return ws
317+
coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
318+
return r.Workspace
319319
}
320320

321321
// setupTestListener starts accepting connections and echoing a single packet.

‎cli/schedule_test.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,30 +38,30 @@ func setupTestSchedule(t *testing.T, sched *cron.Schedule) (ownerClient, memberC
3838
memberClient, memberUser := coderdtest.CreateAnotherUserMutators(t, ownerClient, owner.OrganizationID, nil, func(r *codersdk.CreateUserRequest) {
3939
r.Username = "testuser2" // ensure deterministic ordering
4040
})
41-
_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
41+
_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
4242
Name: "a-owner",
4343
OwnerID: owner.UserID,
4444
OrganizationID: owner.OrganizationID,
4545
AutostartSchedule: sql.NullString{String: sched.String(), Valid: true},
4646
Ttl: sql.NullInt64{Int64: 8 * time.Hour.Nanoseconds(), Valid: true},
47-
})
48-
_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
47+
}).WithAgent().Do()
48+
_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
4949
Name: "b-owner",
5050
OwnerID: owner.UserID,
5151
OrganizationID: owner.OrganizationID,
5252
AutostartSchedule: sql.NullString{String: sched.String(), Valid: true},
53-
})
54-
_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
53+
}).WithAgent().Do()
54+
_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
5555
Name: "c-member",
5656
OwnerID: memberUser.ID,
5757
OrganizationID: owner.OrganizationID,
5858
Ttl: sql.NullInt64{Int64: 8 * time.Hour.Nanoseconds(), Valid: true},
59-
})
60-
_, _ = dbfake.WorkspaceWithAgent(t, db, database.Workspace{
59+
}).WithAgent().Do()
60+
_ = dbfake.NewWorkspaceBuilder(t, db).Seed(database.Workspace{
6161
Name: "d-member",
6262
OwnerID: memberUser.ID,
6363
OrganizationID: owner.OrganizationID,
64-
})
64+
}).WithAgent().Do()
6565

6666
// Need this for LatestBuild.Deadline
6767
resp, err := ownerClient.Workspaces(context.Background(), codersdk.WorkspaceFilter{})

‎cli/ssh_test.go

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -52,12 +52,15 @@ func setupWorkspaceForAgent(t *testing.T, mutations ...func([]*proto.Agent) []*p
5252
client.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
5353
first := coderdtest.CreateFirstUser(t, client)
5454
userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
55-
workspace, agentToken := dbfake.WorkspaceWithAgent(t, store, database.Workspace{
56-
OrganizationID: first.OrganizationID,
57-
OwnerID: user.ID,
58-
}, mutations...)
55+
r := dbfake.NewWorkspaceBuilder(t, store).
56+
Seed(database.Workspace{
57+
OrganizationID: first.OrganizationID,
58+
OwnerID: user.ID,
59+
}).
60+
WithAgent(mutations...).
61+
Do()
5962

60-
return userClient, workspace, agentToken
63+
return userClient, r.Workspace, r.AgentToken
6164
}
6265

6366
func TestSSH(t *testing.T) {
@@ -127,11 +130,11 @@ func TestSSH(t *testing.T) {
127130
client.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
128131
first := coderdtest.CreateFirstUser(t, client)
129132
userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
130-
workspace, agentToken := dbfake.WorkspaceWithAgent(t, store, database.Workspace{
133+
r := dbfake.NewWorkspaceBuilder(t, store).Seed(database.Workspace{
131134
OrganizationID: first.OrganizationID,
132135
OwnerID: user.ID,
133-
})
134-
inv, root := clitest.New(t, "ssh", workspace.Name)
136+
}).WithAgent().Do()
137+
inv, root := clitest.New(t, "ssh", r.Workspace.Name)
135138
clitest.SetupConfig(t, userClient, root)
136139
pty := ptytest.New(t).Attach(inv)
137140

@@ -144,14 +147,14 @@ func TestSSH(t *testing.T) {
144147
})
145148
pty.ExpectMatch("Waiting")
146149

147-
_ = agenttest.New(t, client.URL, agentToken)
148-
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
150+
_ = agenttest.New(t, client.URL, r.AgentToken)
151+
coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
149152

150153
// Ensure the agent is connected.
151154
pty.WriteLine("echo hell'o'")
152155
pty.ExpectMatchContext(ctx, "hello")
153156

154-
_ = dbfake.NewWorkspaceBuildBuilder(t, store, workspace).
157+
_ = dbfake.NewWorkspaceBuildBuilder(t, store, r.Workspace).
155158
Seed(database.WorkspaceBuild{
156159
Transition: database.WorkspaceTransitionStop,
157160
BuildNumber: 2,
@@ -466,15 +469,15 @@ func TestSSH(t *testing.T) {
466469
client.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
467470
first := coderdtest.CreateFirstUser(t, client)
468471
userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
469-
workspace, agentToken := dbfake.WorkspaceWithAgent(t, store, database.Workspace{
472+
r := dbfake.NewWorkspaceBuilder(t, store).Seed(database.Workspace{
470473
OrganizationID: first.OrganizationID,
471474
OwnerID: user.ID,
472-
})
475+
}).WithAgent().Do()
473476

474477
_, _ = tGoContext(t, func(ctx context.Context) {
475478
// Run this async so the SSH command has to wait for
476479
// the build and agent to connect.
477-
_ = agenttest.New(t, client.URL, agentToken)
480+
_ = agenttest.New(t, client.URL, r.AgentToken)
478481
<-ctx.Done()
479482
})
480483

@@ -489,7 +492,7 @@ func TestSSH(t *testing.T) {
489492
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
490493
defer cancel()
491494

492-
inv, root := clitest.New(t, "ssh", "--stdio", workspace.Name)
495+
inv, root := clitest.New(t, "ssh", "--stdio", r.Workspace.Name)
493496
clitest.SetupConfig(t, userClient, root)
494497
inv.Stdin = clientOutput
495498
inv.Stdout = serverInput
@@ -520,7 +523,7 @@ func TestSSH(t *testing.T) {
520523
err = session.Shell()
521524
require.NoError(t, err)
522525

523-
_ = dbfake.NewWorkspaceBuildBuilder(t, store, workspace).
526+
_ = dbfake.NewWorkspaceBuildBuilder(t, store, r.Workspace).
524527
Seed(database.WorkspaceBuild{
525528
Transition: database.WorkspaceTransitionStop,
526529
BuildNumber: 2,

‎coderd/database/dbfake/dbfake.go

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ func (b WorkspaceBuilder) Seed(seed database.Workspace) WorkspaceBuilder {
4747
}
4848

4949
func (b WorkspaceBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) []*sdkproto.Agent) WorkspaceBuilder {
50+
//nolint: revive // returns modified struct
5051
b.agentToken = uuid.NewString()
5152
agents := []*sdkproto.Agent{{
5253
Id: uuid.NewString(),
@@ -95,19 +96,6 @@ func Workspace(t testing.TB, db database.Store, seed database.Workspace) databas
9596
return r.Workspace
9697
}
9798

98-
// WorkspaceWithAgent is a helper that generates a workspace with a single resource
99-
// that has an agent attached to it. The agent token is returned.
100-
func WorkspaceWithAgent(
101-
t testing.TB, db database.Store, seed database.Workspace,
102-
mutations ...func([]*sdkproto.Agent) []*sdkproto.Agent,
103-
) (
104-
database.Workspace, string,
105-
) {
106-
t.Helper()
107-
r := NewWorkspaceBuilder(t, db).Seed(seed).WithAgent(mutations...).Do()
108-
return r.Workspace, r.AgentToken
109-
}
110-
11199
type WorkspaceBuildBuilder struct {
112100
t testing.TB
113101
db database.Store

0 commit comments

Comments
 (0)