Skip to content

Commit 93ef696

Browse files
johnstcnjohnstcn
authored
refactor(agent): add agenttest.New helper function (#9812)
* Adds agenttest.New() helper function * Makes sure agent gets closed on test cleanup * Makes sure you don't forget to set session token * Sets the agent and client logger automatically
1 parent 96c5076 commit 93ef696

22 files changed

+153
-368
lines changed

‎agent/agenttest/agent.go

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package agenttest
2+
3+
import (
4+
"context"
5+
"net/url"
6+
"testing"
7+
8+
"github.com/stretchr/testify/assert"
9+
10+
"cdr.dev/slog"
11+
"cdr.dev/slog/sloggers/slogtest"
12+
"github.com/coder/coder/v2/agent"
13+
"github.com/coder/coder/v2/codersdk/agentsdk"
14+
)
15+
16+
// New starts a new agent for use in tests.
17+
// The agent will use the provided coder URL and session token.
18+
// The options passed to agent.New() can be modified by passing an optional
19+
// variadic func(*agent.Options).
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 {
24+
t.Helper()
25+
26+
var o agent.Options
27+
log := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent")
28+
o.Logger = log
29+
30+
for _, opt := range opts {
31+
opt(&o)
32+
}
33+
34+
if o.Client == nil {
35+
agentClient := agentsdk.New(coderURL)
36+
agentClient.SetSessionToken(agentToken)
37+
agentClient.SDK.SetLogger(log)
38+
o.Client = agentClient
39+
}
40+
41+
if o.ExchangeToken == nil {
42+
o.ExchangeToken = func(_ context.Context) (string, error) {
43+
return agentToken, nil
44+
}
45+
}
46+
47+
if o.LogDir == "" {
48+
o.LogDir = t.TempDir()
49+
}
50+
51+
agt := agent.New(o)
52+
t.Cleanup(func() {
53+
assert.NoError(t, agt.Close(), "failed to close agent during cleanup")
54+
})
55+
56+
return agt
57+
}

‎cli/configssh_test.go

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,10 @@ import (
1919
"github.com/stretchr/testify/assert"
2020
"github.com/stretchr/testify/require"
2121

22-
"cdr.dev/slog/sloggers/slogtest"
23-
24-
"github.com/coder/coder/v2/agent"
22+
"github.com/coder/coder/v2/agent/agenttest"
2523
"github.com/coder/coder/v2/cli/clitest"
2624
"github.com/coder/coder/v2/coderd/coderdtest"
2725
"github.com/coder/coder/v2/codersdk"
28-
"github.com/coder/coder/v2/codersdk/agentsdk"
2926
"github.com/coder/coder/v2/provisioner/echo"
3027
"github.com/coder/coder/v2/provisionersdk/proto"
3128
"github.com/coder/coder/v2/pty/ptytest"
@@ -102,15 +99,7 @@ func TestConfigSSH(t *testing.T) {
10299
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
103100
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
104101
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
105-
agentClient := agentsdk.New(client.URL)
106-
agentClient.SetSessionToken(authToken)
107-
agentCloser := agent.New(agent.Options{
108-
Client: agentClient,
109-
Logger: slogtest.Make(t, nil).Named("agent"),
110-
})
111-
defer func() {
112-
_ = agentCloser.Close()
113-
}()
102+
_ = agenttest.New(t, client.URL, authToken)
114103
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
115104
agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
116105
require.NoError(t, err)

‎cli/gitssh_test.go

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

23+
"github.com/coder/coder/v2/agent"
24+
"github.com/coder/coder/v2/agent/agenttest"
2325
"github.com/coder/coder/v2/cli/clitest"
2426
"github.com/coder/coder/v2/coderd/coderdtest"
2527
"github.com/coder/coder/v2/codersdk"
28+
"github.com/coder/coder/v2/codersdk/agentsdk"
2629
"github.com/coder/coder/v2/provisioner/echo"
2730
"github.com/coder/coder/v2/pty/ptytest"
2831
"github.com/coder/coder/v2/testutil"
2932
)
3033

31-
func prepareTestGitSSH(ctx context.Context, t *testing.T) (*codersdk.Client, string, gossh.PublicKey) {
34+
func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, string, gossh.PublicKey) {
3235
t.Helper()
3336

3437
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
@@ -57,13 +60,12 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*codersdk.Client, str
5760
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
5861

5962
// start workspace agent
60-
inv, root := clitest.New(t, "agent", "--agent-token", agentToken, "--agent-url", client.URL.String())
61-
agentClient := codersdk.New(client.URL)
63+
agentClient := agentsdk.New(client.URL)
6264
agentClient.SetSessionToken(agentToken)
63-
clitest.SetupConfig(t, agentClient, root)
64-
clitest.Start(t, inv)
65-
66-
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
65+
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
66+
o.Client = agentClient
67+
})
68+
_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
6769
return agentClient, agentToken, pubkey
6870
}
6971

@@ -140,7 +142,7 @@ func TestGitSSH(t *testing.T) {
140142
// set to agent config dir
141143
inv, _ := clitest.New(t,
142144
"gitssh",
143-
"--agent-url", client.URL.String(),
145+
"--agent-url", client.SDK.URL.String(),
144146
"--agent-token", token,
145147
"--",
146148
fmt.Sprintf("-p%d", addr.Port),
@@ -203,7 +205,7 @@ func TestGitSSH(t *testing.T) {
203205
pty := ptytest.New(t)
204206
cmdArgs := []string{
205207
"gitssh",
206-
"--agent-url", client.URL.String(),
208+
"--agent-url", client.SDK.URL.String(),
207209
"--agent-token", token,
208210
"--",
209211
"-F", config,

‎cli/ping_test.go

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,9 @@ import (
66

77
"github.com/stretchr/testify/assert"
88

9-
"cdr.dev/slog/sloggers/slogtest"
10-
11-
"github.com/coder/coder/v2/agent"
9+
"github.com/coder/coder/v2/agent/agenttest"
1210
"github.com/coder/coder/v2/cli/clitest"
13-
"github.com/coder/coder/v2/codersdk/agentsdk"
11+
"github.com/coder/coder/v2/coderd/coderdtest"
1412
"github.com/coder/coder/v2/pty/ptytest"
1513
"github.com/coder/coder/v2/testutil"
1614
)
@@ -29,15 +27,8 @@ func TestPing(t *testing.T) {
2927
inv.Stderr = pty.Output()
3028
inv.Stdout = pty.Output()
3129

32-
agentClient := agentsdk.New(client.URL)
33-
agentClient.SetSessionToken(agentToken)
34-
agentCloser := agent.New(agent.Options{
35-
Client: agentClient,
36-
Logger: slogtest.Make(t, nil).Named("agent"),
37-
})
38-
defer func() {
39-
_ = agentCloser.Close()
40-
}()
30+
_ = agenttest.New(t, client.URL, agentToken)
31+
_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
4132

4233
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
4334
defer cancel()

‎cli/portforward_test.go

Lines changed: 5 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,11 @@ import (
1414
"github.com/stretchr/testify/assert"
1515
"github.com/stretchr/testify/require"
1616

17-
"cdr.dev/slog"
18-
"cdr.dev/slog/sloggers/slogtest"
1917
"github.com/coder/coder/v2/agent"
18+
"github.com/coder/coder/v2/agent/agenttest"
2019
"github.com/coder/coder/v2/cli/clitest"
2120
"github.com/coder/coder/v2/coderd/coderdtest"
2221
"github.com/coder/coder/v2/codersdk"
23-
"github.com/coder/coder/v2/codersdk/agentsdk"
2422
"github.com/coder/coder/v2/provisioner/echo"
2523
"github.com/coder/coder/v2/pty/ptytest"
2624
"github.com/coder/coder/v2/testutil"
@@ -317,23 +315,11 @@ func runAgent(t *testing.T, client *codersdk.Client, userID uuid.UUID) codersdk.
317315
workspace := coderdtest.CreateWorkspace(t, client, orgID, template.ID)
318316
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
319317

320-
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent")
321-
agentClient := agentsdk.New(client.URL)
322-
agentClient.SDK.SetLogger(logger)
323-
agentClient.SDK.SetSessionToken(agentToken)
324-
agnt := agent.New(agent.Options{
325-
Client: agentClient,
326-
Logger: logger,
327-
LogDir: t.TempDir(),
328-
ExchangeToken: func(ctx context.Context) (string, error) {
329-
return agentToken, nil
318+
_ = agenttest.New(t, client.URL, agentToken,
319+
func(o *agent.Options) {
320+
o.SSHMaxTimeout = 60 * time.Second
330321
},
331-
SSHMaxTimeout: time.Second * 60,
332-
})
333-
t.Cleanup(func() {
334-
err := agnt.Close()
335-
assert.NoError(t, err)
336-
})
322+
)
337323
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
338324

339325
return workspace

‎cli/speedtest_test.go

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,11 @@ import (
99

1010
"cdr.dev/slog"
1111
"cdr.dev/slog/sloggers/slogtest"
12-
"github.com/coder/coder/v2/agent"
12+
"github.com/coder/coder/v2/agent/agenttest"
1313
"github.com/coder/coder/v2/cli"
1414
"github.com/coder/coder/v2/cli/clitest"
1515
"github.com/coder/coder/v2/coderd/coderdtest"
1616
"github.com/coder/coder/v2/codersdk"
17-
"github.com/coder/coder/v2/codersdk/agentsdk"
1817
"github.com/coder/coder/v2/pty/ptytest"
1918
"github.com/coder/coder/v2/testutil"
2019
)
@@ -26,13 +25,7 @@ func TestSpeedtest(t *testing.T) {
2625
t.Skip("This test takes a minimum of 5ms per a hardcoded value in Tailscale!")
2726
}
2827
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
29-
agentClient := agentsdk.New(client.URL)
30-
agentClient.SetSessionToken(agentToken)
31-
agentCloser := agent.New(agent.Options{
32-
Client: agentClient,
33-
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
34-
})
35-
defer agentCloser.Close()
28+
_ = agenttest.New(t, client.URL, agentToken)
3629
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
3730

3831
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

0 commit comments

Comments
 (0)