Skip to content

refactor(agent): add agenttest.New helper function #9812

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Sep 26, 2023
97 changes: 97 additions & 0 deletions agent/agenttest/agent.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
package agenttest

import (
"context"
"net/url"
"sync"
"testing"

"github.com/google/uuid"
"github.com/stretchr/testify/assert"

"cdr.dev/slog"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
)

// Agent is a small wrapper around an agent for use in tests.
type Agent struct {
agent agent.Agent
agentClient *agentsdk.Client
resources []codersdk.WorkspaceResource
waiter func(*codersdk.Client, uuid.UUID, ...string) []codersdk.WorkspaceResource
waitOnce sync.Once
}

// Wait waits for the agent to connect to the workspace and returns the
// resources for the connected workspace.
// Calls coderdtest.AwaitWorkspaceAgents under the hood.
// Multiple calls to Wait() are idempotent.
func (a *Agent) Wait(client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
a.waitOnce.Do(func() {
a.resources = a.waiter(client, workspaceID, agentNames...)
})
return a.resources
}

// Client returns the agent client.
func (a *Agent) Client() *agentsdk.Client {
return a.agentClient
}

// Agent returns the agent itself.
func (a *Agent) Agent() agent.Agent {
return a.agent
}

// New starts a new agent for use in tests.
// The agent will use the provided coder URL and session token.
// The options passed to agent.New() can be modified by passing an optional
// variadic func(*agent.Options).
// Returns a wrapper that can be used to wait for the agent to connect to the
// workspace by calling Wait(). The arguments to Wait() are passed to
// coderdtest.AwaitWorkspaceAgents.
// Closing the agent is handled by the test cleanup.
func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent.Options)) *Agent {
t.Helper()

var o agent.Options
log := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent")
o.Logger = log

for _, opt := range opts {
opt(&o)
}

if o.Client == nil {
agentClient := agentsdk.New(coderURL)
agentClient.SetSessionToken(agentToken)
o.Client = agentClient
}

if o.ExchangeToken == nil {
o.ExchangeToken = func(_ context.Context) (string, error) {
return agentToken, nil
}
}

if o.LogDir == "" {
o.LogDir = t.TempDir()
}

agentCloser := agent.New(o)
t.Cleanup(func() {
assert.NoError(t, agentCloser.Close(), "failed to close agent during cleanup")
})

return &Agent{
agent: agentCloser,
agentClient: o.Client.(*agentsdk.Client), // nolint:forcetypeassert
waiter: func(c *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
return coderdtest.AwaitWorkspaceAgents(t, c, workspaceID, agentNames...)
},
}
}
16 changes: 2 additions & 14 deletions cli/configssh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,10 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"cdr.dev/slog/sloggers/slogtest"

"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/agent/agenttest"
"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/provisioner/echo"
"github.com/coder/coder/v2/provisionersdk/proto"
"github.com/coder/coder/v2/pty/ptytest"
Expand Down Expand Up @@ -102,16 +99,7 @@ func TestConfigSSH(t *testing.T) {
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
agentClient := agentsdk.New(client.URL)
agentClient.SetSessionToken(authToken)
agentCloser := agent.New(agent.Options{
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent"),
})
defer func() {
_ = agentCloser.Close()
}()
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID)
agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
require.NoError(t, err)
defer agentConn.Close()
Expand Down
20 changes: 9 additions & 11 deletions cli/gitssh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ import (
"github.com/stretchr/testify/require"
gossh "golang.org/x/crypto/ssh"

"github.com/coder/coder/v2/agent/agenttest"
"github.com/coder/coder/v2/codersdk/agentsdk"

"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/codersdk"
Expand All @@ -28,7 +31,7 @@ import (
"github.com/coder/coder/v2/testutil"
)

func prepareTestGitSSH(ctx context.Context, t *testing.T) (*codersdk.Client, string, gossh.PublicKey) {
func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, string, gossh.PublicKey) {
t.Helper()

client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
Expand Down Expand Up @@ -57,14 +60,9 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*codersdk.Client, str
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)

// start workspace agent
inv, root := clitest.New(t, "agent", "--agent-token", agentToken, "--agent-url", client.URL.String())
agentClient := codersdk.New(client.URL)
agentClient.SetSessionToken(agentToken)
clitest.SetupConfig(t, agentClient, root)
clitest.Start(t, inv)

coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
return agentClient, agentToken, pubkey
agt := agenttest.New(t, client.URL, agentToken)
agt.Wait(client, workspace.ID)
return agt.Client(), agentToken, pubkey
}

func serveSSHForGitSSH(t *testing.T, handler func(ssh.Session), pubkeys ...gossh.PublicKey) *net.TCPAddr {
Expand Down Expand Up @@ -140,7 +138,7 @@ func TestGitSSH(t *testing.T) {
// set to agent config dir
inv, _ := clitest.New(t,
"gitssh",
"--agent-url", client.URL.String(),
"--agent-url", client.SDK.URL.String(),
"--agent-token", token,
"--",
fmt.Sprintf("-p%d", addr.Port),
Expand Down Expand Up @@ -203,7 +201,7 @@ func TestGitSSH(t *testing.T) {
pty := ptytest.New(t)
cmdArgs := []string{
"gitssh",
"--agent-url", client.URL.String(),
"--agent-url", client.SDK.URL.String(),
"--agent-token", token,
"--",
"-F", config,
Expand Down
15 changes: 2 additions & 13 deletions cli/ping_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,8 @@ import (

"github.com/stretchr/testify/assert"

"cdr.dev/slog/sloggers/slogtest"

"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/agent/agenttest"
"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/pty/ptytest"
"github.com/coder/coder/v2/testutil"
)
Expand All @@ -29,15 +26,7 @@ func TestPing(t *testing.T) {
inv.Stderr = pty.Output()
inv.Stdout = pty.Output()

agentClient := agentsdk.New(client.URL)
agentClient.SetSessionToken(agentToken)
agentCloser := agent.New(agent.Options{
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent"),
})
defer func() {
_ = agentCloser.Close()
}()
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
Expand Down
26 changes: 6 additions & 20 deletions cli/portforward_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,12 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"cdr.dev/slog"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/agent/agenttest"

"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/provisioner/echo"
"github.com/coder/coder/v2/pty/ptytest"
"github.com/coder/coder/v2/testutil"
Expand Down Expand Up @@ -317,24 +316,11 @@ func runAgent(t *testing.T, client *codersdk.Client, userID uuid.UUID) codersdk.
workspace := coderdtest.CreateWorkspace(t, client, orgID, template.ID)
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)

logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent")
agentClient := agentsdk.New(client.URL)
agentClient.SDK.SetLogger(logger)
agentClient.SDK.SetSessionToken(agentToken)
agnt := agent.New(agent.Options{
Client: agentClient,
Logger: logger,
LogDir: t.TempDir(),
ExchangeToken: func(ctx context.Context) (string, error) {
return agentToken, nil
agenttest.New(t, client.URL, agentToken,
func(o *agent.Options) {
o.SSHMaxTimeout = 60 * time.Second
},
SSHMaxTimeout: time.Second * 60,
})
t.Cleanup(func() {
err := agnt.Close()
assert.NoError(t, err)
})
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
).Wait(client, workspace.ID)

return workspace
}
Expand Down
13 changes: 2 additions & 11 deletions cli/speedtest_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,10 @@ import (

"cdr.dev/slog"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/agent/agenttest"
"github.com/coder/coder/v2/cli"
"github.com/coder/coder/v2/cli/clitest"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/pty/ptytest"
"github.com/coder/coder/v2/testutil"
)
Expand All @@ -26,14 +24,7 @@ func TestSpeedtest(t *testing.T) {
t.Skip("This test takes a minimum of 5ms per a hardcoded value in Tailscale!")
}
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
agentClient := agentsdk.New(client.URL)
agentClient.SetSessionToken(agentToken)
agentCloser := agent.New(agent.Options{
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
})
defer agentCloser.Close()
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
_ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID)

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
Expand Down
Loading