Skip to content

Commit d55c797

Browse files
johnstcnjohnstcn
committed
address PR comments
1 parent 6aa4bbe commit d55c797

20 files changed

+89
-96
lines changed

‎agent/agenttest/agent.go

Lines changed: 8 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -3,59 +3,24 @@ package agenttest
33
import (
44
"context"
55
"net/url"
6-
"sync"
76
"testing"
87

9-
"github.com/google/uuid"
108
"github.com/stretchr/testify/assert"
119

1210
"cdr.dev/slog"
1311
"cdr.dev/slog/sloggers/slogtest"
1412
"github.com/coder/coder/v2/agent"
15-
"github.com/coder/coder/v2/coderd/coderdtest"
16-
"github.com/coder/coder/v2/codersdk"
1713
"github.com/coder/coder/v2/codersdk/agentsdk"
1814
)
1915

20-
// Agent is a small wrapper around an agent for use in tests.
21-
type Agent struct {
22-
agent agent.Agent
23-
agentClient *agentsdk.Client
24-
resources []codersdk.WorkspaceResource
25-
waiter func(*codersdk.Client, uuid.UUID, ...string) []codersdk.WorkspaceResource
26-
waitOnce sync.Once
27-
}
28-
29-
// Wait waits for the agent to connect to the workspace and returns the
30-
// resources for the connected workspace.
31-
// Calls coderdtest.AwaitWorkspaceAgents under the hood.
32-
// Multiple calls to Wait() are idempotent.
33-
func (a *Agent) Wait(client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
34-
a.waitOnce.Do(func() {
35-
a.resources = a.waiter(client, workspaceID, agentNames...)
36-
})
37-
return a.resources
38-
}
39-
40-
// Client returns the agent client.
41-
func (a *Agent) Client() *agentsdk.Client {
42-
return a.agentClient
43-
}
44-
45-
// Agent returns the agent itself.
46-
func (a *Agent) Agent() agent.Agent {
47-
return a.agent
48-
}
49-
5016
// New starts a new agent for use in tests.
5117
// The agent will use the provided coder URL and session token.
5218
// The options passed to agent.New() can be modified by passing an optional
5319
// variadic func(*agent.Options).
54-
// Returns a wrapper that can be used to wait for the agent to connect to the
55-
// workspace by calling Wait(). The arguments to Wait() are passed to
56-
// coderdtest.AwaitWorkspaceAgents.
57-
// Closing the agent is handled by the test cleanup.
58-
func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent.Options)) *Agent {
20+
// Returns the agent. Closing the agent is handled by the test cleanup.
21+
// It is the responsibility of the caller to call coderdtest.AwaitWorkspaceAgents
22+
// to ensure agent is connected.
23+
func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent.Options)) agent.Agent {
5924
t.Helper()
6025

6126
var o agent.Options
@@ -69,6 +34,7 @@ func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent
6934
if o.Client == nil {
7035
agentClient := agentsdk.New(coderURL)
7136
agentClient.SetSessionToken(agentToken)
37+
agentClient.SDK.SetLogger(log)
7238
o.Client = agentClient
7339
}
7440

@@ -82,16 +48,10 @@ func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent
8248
o.LogDir = t.TempDir()
8349
}
8450

85-
agentCloser := agent.New(o)
51+
agt := agent.New(o)
8652
t.Cleanup(func() {
87-
assert.NoError(t, agentCloser.Close(), "failed to close agent during cleanup")
53+
assert.NoError(t, agt.Close(), "failed to close agent during cleanup")
8854
})
8955

90-
return &Agent{
91-
agent: agentCloser,
92-
agentClient: o.Client.(*agentsdk.Client), // nolint:forcetypeassert
93-
waiter: func(c *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
94-
return coderdtest.AwaitWorkspaceAgents(t, c, workspaceID, agentNames...)
95-
},
96-
}
56+
return agt
9757
}

‎cli/configssh_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,8 @@ func TestConfigSSH(t *testing.T) {
9999
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
100100
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
101101
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
102-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
102+
_ = agenttest.New(t, client.URL, authToken)
103+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
103104
agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
104105
require.NoError(t, err)
105106
defer agentConn.Close()

‎cli/gitssh_test.go

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,12 @@ import (
2020
"github.com/stretchr/testify/require"
2121
gossh "golang.org/x/crypto/ssh"
2222

23+
"github.com/coder/coder/v2/agent"
2324
"github.com/coder/coder/v2/agent/agenttest"
24-
"github.com/coder/coder/v2/codersdk/agentsdk"
25-
2625
"github.com/coder/coder/v2/cli/clitest"
2726
"github.com/coder/coder/v2/coderd/coderdtest"
2827
"github.com/coder/coder/v2/codersdk"
28+
"github.com/coder/coder/v2/codersdk/agentsdk"
2929
"github.com/coder/coder/v2/provisioner/echo"
3030
"github.com/coder/coder/v2/pty/ptytest"
3131
"github.com/coder/coder/v2/testutil"
@@ -60,9 +60,13 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str
6060
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
6161

6262
// start workspace agent
63-
agt := agenttest.New(t, client.URL, agentToken)
64-
agt.Wait(client, workspace.ID)
65-
return agt.Client(), agentToken, pubkey
63+
agentClient := agentsdk.New(client.URL)
64+
agentClient.SetSessionToken(agentToken)
65+
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
66+
o.Client = agentClient
67+
})
68+
_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
69+
return agentClient, agentToken, pubkey
6670
}
6771

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

‎cli/ping_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
"github.com/coder/coder/v2/agent/agenttest"
1010
"github.com/coder/coder/v2/cli/clitest"
11+
"github.com/coder/coder/v2/coderd/coderdtest"
1112
"github.com/coder/coder/v2/pty/ptytest"
1213
"github.com/coder/coder/v2/testutil"
1314
)
@@ -26,7 +27,8 @@ func TestPing(t *testing.T) {
2627
inv.Stderr = pty.Output()
2728
inv.Stdout = pty.Output()
2829

29-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
30+
_ = agenttest.New(t, client.URL, agentToken)
31+
_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
3032

3133
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
3234
defer cancel()

‎cli/portforward_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import (
1616

1717
"github.com/coder/coder/v2/agent"
1818
"github.com/coder/coder/v2/agent/agenttest"
19-
2019
"github.com/coder/coder/v2/cli/clitest"
2120
"github.com/coder/coder/v2/coderd/coderdtest"
2221
"github.com/coder/coder/v2/codersdk"
@@ -316,11 +315,12 @@ func runAgent(t *testing.T, client *codersdk.Client, userID uuid.UUID) codersdk.
316315
workspace := coderdtest.CreateWorkspace(t, client, orgID, template.ID)
317316
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
318317

319-
agenttest.New(t, client.URL, agentToken,
318+
_ = agenttest.New(t, client.URL, agentToken,
320319
func(o *agent.Options) {
321320
o.SSHMaxTimeout = 60 * time.Second
322321
},
323-
).Wait(client, workspace.ID)
322+
)
323+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
324324

325325
return workspace
326326
}

‎cli/speedtest_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/coder/coder/v2/agent/agenttest"
1313
"github.com/coder/coder/v2/cli"
1414
"github.com/coder/coder/v2/cli/clitest"
15+
"github.com/coder/coder/v2/coderd/coderdtest"
1516
"github.com/coder/coder/v2/codersdk"
1617
"github.com/coder/coder/v2/pty/ptytest"
1718
"github.com/coder/coder/v2/testutil"
@@ -24,7 +25,8 @@ func TestSpeedtest(t *testing.T) {
2425
t.Skip("This test takes a minimum of 5ms per a hardcoded value in Tailscale!")
2526
}
2627
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
27-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
28+
_ = agenttest.New(t, client.URL, agentToken)
29+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
2830

2931
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
3032
defer cancel()

‎cli/ssh_test.go

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,8 @@ func TestSSH(t *testing.T) {
103103
})
104104
pty.ExpectMatch("Waiting")
105105

106-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
106+
_ = agenttest.New(t, client.URL, agentToken)
107+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
107108

108109
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
109110
pty.WriteLine("exit")
@@ -159,7 +160,8 @@ func TestSSH(t *testing.T) {
159160
})
160161
pty.ExpectMatch("Waiting")
161162

162-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
163+
_ = agenttest.New(t, client.URL, agentToken)
164+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
163165

164166
// Ensure the agent is connected.
165167
pty.WriteLine("echo hell'o'")
@@ -180,7 +182,8 @@ func TestSSH(t *testing.T) {
180182
_, _ = tGoContext(t, func(ctx context.Context) {
181183
// Run this async so the SSH command has to wait for
182184
// the build and agent to connect!
183-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
185+
_ = agenttest.New(t, client.URL, agentToken)
186+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
184187
<-ctx.Done()
185188
})
186189

@@ -242,7 +245,8 @@ func TestSSH(t *testing.T) {
242245
_, _ = tGoContext(t, func(ctx context.Context) {
243246
// Run this async so the SSH command has to wait for
244247
// the build and agent to connect.
245-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
248+
_ = agenttest.New(t, client.URL, agentToken)
249+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
246250
<-ctx.Done()
247251
})
248252

@@ -305,7 +309,8 @@ func TestSSH(t *testing.T) {
305309

306310
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
307311

308-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
312+
_ = agenttest.New(t, client.URL, agentToken)
313+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
309314

310315
// Generate private key.
311316
privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
@@ -390,7 +395,8 @@ func TestSSH(t *testing.T) {
390395

391396
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
392397

393-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
398+
_ = agenttest.New(t, client.URL, agentToken)
399+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
394400

395401
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
396402
defer cancel()
@@ -435,7 +441,8 @@ func TestSSH(t *testing.T) {
435441

436442
pty.ExpectMatch("Waiting")
437443

438-
agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
444+
agenttest.New(t, client.URL, agentToken)
445+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
439446

440447
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
441448
pty.WriteLine("exit")
@@ -607,8 +614,8 @@ Expire-Date: 0
607614
o.EnvironmentVariables = map[string]string{
608615
"GNUPGHOME": gnupgHomeWorkspace,
609616
}
610-
},
611-
).Wait(client, workspace.ID)
617+
})
618+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
612619

613620
inv, root := clitest.New(t,
614621
"ssh",

‎cli/vscodessh_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111

1212
"github.com/coder/coder/v2/agent/agenttest"
1313
"github.com/coder/coder/v2/cli/clitest"
14+
"github.com/coder/coder/v2/coderd/coderdtest"
1415
"github.com/coder/coder/v2/codersdk"
1516
"github.com/coder/coder/v2/pty/ptytest"
1617
"github.com/coder/coder/v2/testutil"
@@ -25,7 +26,8 @@ func TestVSCodeSSH(t *testing.T) {
2526
user, err := client.User(ctx, codersdk.Me)
2627
require.NoError(t, err)
2728

28-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
29+
_ = agenttest.New(t, client.URL, agentToken)
30+
_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
2931

3032
fs := afero.NewMemMapFs()
3133
err = afero.WriteFile(fs, "/url", []byte(client.URL.String()), 0o600)

‎coderd/activitybump_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,8 @@ func TestWorkspaceActivityBump(t *testing.T) {
8585
require.NoError(t, err)
8686
}
8787

88-
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)
88+
_ = agenttest.New(t, client.URL, agentToken)
89+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
8990

9091
// Sanity-check that deadline is near.
9192
workspace, err := client.Workspace(ctx, workspace.ID)

‎coderd/coderd_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,8 @@ func TestDERPForceWebSockets(t *testing.T) {
188188
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
189189
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
190190

191-
_ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
191+
_ = agenttest.New(t, client.URL, authToken)
192+
_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
192193

193194
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
194195
defer cancel()

‎coderd/insights_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ func TestDeploymentInsights(t *testing.T) {
5858
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
5959
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
6060

61-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
61+
_ = agenttest.New(t, client.URL, authToken)
62+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
6263
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
6364
defer cancel()
6465

@@ -135,7 +136,8 @@ func TestUserLatencyInsights(t *testing.T) {
135136
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
136137

137138
// Start an agent so that we can generate stats.
138-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
139+
_ = agenttest.New(t, client.URL, authToken)
140+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
139141

140142
// Start must be at the beginning of the day, initialize it early in case
141143
// the day changes so that we get the relevant stats faster.

‎coderd/templates_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1217,7 +1217,8 @@ func TestTemplateMetrics(t *testing.T) {
12171217
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
12181218
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
12191219

1220-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
1220+
_ = agenttest.New(t, client.URL, authToken)
1221+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
12211222

12221223
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
12231224
defer cancel()

‎coderd/workspaceagents_test.go

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -484,7 +484,8 @@ func TestWorkspaceAgentListen(t *testing.T) {
484484
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
485485
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
486486

487-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
487+
_ = agenttest.New(t, client.URL, authToken)
488+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
488489

489490
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
490491
defer cancel()
@@ -576,7 +577,8 @@ func TestWorkspaceAgentTailnet(t *testing.T) {
576577
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
577578
daemonCloser.Close()
578579

579-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
580+
_ = agenttest.New(t, client.URL, authToken)
581+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
580582

581583
ctx, cancelFunc := context.WithCancel(context.Background())
582584
defer cancelFunc()
@@ -630,7 +632,8 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) {
630632
require.NoError(t, err)
631633
require.True(t, manifest.DisableDirectConnections)
632634

633-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
635+
_ = agenttest.New(t, client.URL, authToken)
636+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
634637
agentID := resources[0].Agents[0].ID
635638

636639
// Verify that the connection data has no STUN ports and
@@ -708,7 +711,8 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) {
708711
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
709712
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
710713

711-
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
714+
_ = agenttest.New(t, client.URL, authToken)
715+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
712716
return client, uint16(coderdPort), resources[0].Agents[0].ID
713717
}
714718

@@ -1427,10 +1431,9 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) {
14271431
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
14281432
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
14291433

1430-
agt := agenttest.New(t, client.URL, agentToken)
1431-
resources := agt.Wait(client, workspace.ID)
1434+
agentCloser := agenttest.New(t, client.URL, agentToken)
1435+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
14321436
agentID := resources[0].Agents[0].ID
1433-
agentCloser := agt.Agent()
14341437

14351438
// Connect from a client.
14361439
ctx := testutil.Context(t, testutil.WaitLong)

0 commit comments

Comments
 (0)