From af64f27e888b523dfb5cf42f050592d4eb8d6a33 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 09:28:35 +0000 Subject: [PATCH 1/9] fix(cli): call agent directly in gitssh tests --- cli/gitssh_test.go | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/cli/gitssh_test.go b/cli/gitssh_test.go index 3e5045acf0288..0241b54b4f2af 100644 --- a/cli/gitssh_test.go +++ b/cli/gitssh_test.go @@ -15,6 +15,13 @@ import ( "sync/atomic" "testing" + "cdr.dev/slog" + "cdr.dev/slog/sloggers/slogtest" + + "github.com/coder/coder/v2/agent" + "github.com/coder/coder/v2/codersdk/agentsdk" + "github.com/stretchr/testify/assert" + "github.com/gliderlabs/ssh" "github.com/google/uuid" "github.com/stretchr/testify/require" @@ -28,7 +35,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}) @@ -57,12 +64,20 @@ 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) + agentLog := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent") + agentClient := agentsdk.New(client.URL) agentClient.SetSessionToken(agentToken) - clitest.SetupConfig(t, agentClient, root) - clitest.Start(t, inv) - + agt := agent.New(agent.Options{ + Client: agentClient, + Logger: agentLog, + LogDir: t.TempDir(), + ExchangeToken: func(_ context.Context) (string, error) { + return agentToken, nil + }, + }) + t.Cleanup(func() { + assert.NoError(t, agt.Close(), "failed to close agent in cleanup") + }) coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) return agentClient, agentToken, pubkey } @@ -140,7 +155,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), @@ -203,7 +218,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, From d99e29c5157facbda4a202288cdbf6c625b8606f Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 12:26:44 +0000 Subject: [PATCH 2/9] add agenttest.New test helper --- agent/agenttest/agent.go | 115 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 agent/agenttest/agent.go diff --git a/agent/agenttest/agent.go b/agent/agenttest/agent.go new file mode 100644 index 0000000000000..fc2086e562de8 --- /dev/null +++ b/agent/agenttest/agent.go @@ -0,0 +1,115 @@ +package agenttest + +import ( + "context" + "net/url" + "testing" + + "github.com/stretchr/testify/require" + + "cdr.dev/slog" + "cdr.dev/slog/sloggers/slogtest" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + + "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" +) + +// Options are options for creating a new test agent. +type Options struct { + // AgentOptions are the options to use for the agent. + AgentOptions agent.Options + + // AgentToken is the token to use for the agent. + AgentToken string + // URL is the URL to which the agent should connect. + URL *url.URL + // WorkspaceID is the ID of the workspace to which the agent should connect. + WorkspaceID uuid.UUID + // Logger is the logger to use for the agent. + // Defaults to a new test logger if not specified. + Logger *slog.Logger +} + +// OptFunc is a function that modifies the given options. +type OptFunc func(*Options) + +func WithAgentToken(token string) OptFunc { + return func(o *Options) { + o.AgentToken = token + } +} + +func WithURL(u *url.URL) OptFunc { + return func(o *Options) { + o.URL = u + } +} + +func WithWorkspaceID(id uuid.UUID) OptFunc { + return func(o *Options) { + o.WorkspaceID = id + } +} + +// New starts a new agent for use in tests. +// Returns the agent client and a function that will block until the agent is +// connected to the workspace. +// Closing the agent is handled by the test cleanup. +func New(t testing.TB, opts ...OptFunc) (agentClient *agentsdk.Client, awaitAgent func(*codersdk.Client)) { + t.Helper() + + var o Options + for _, opt := range opts { + opt(&o) + } + + if o.URL == nil { + require.Fail(t, "must specify URL for agent") + } + agentClient = agentsdk.New(o.URL) + + if o.AgentToken == "" { + o.AgentToken = uuid.NewString() + } + agentClient.SetSessionToken(o.AgentToken) + + if o.AgentOptions.Client == nil { + o.AgentOptions.Client = agentClient + } + + if o.AgentOptions.ExchangeToken == nil { + o.AgentOptions.ExchangeToken = func(_ context.Context) (string, error) { + return o.AgentToken, nil + } + } + + if o.AgentOptions.LogDir == "" { + o.AgentOptions.LogDir = t.TempDir() + } + + if o.Logger == nil { + log := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent") + o.Logger = &log + } + + o.AgentOptions.Logger = *o.Logger + + agentCloser := agent.New(o.AgentOptions) + t.Cleanup(func() { + assert.NoError(t, agentCloser.Close(), "failed to close agent during cleanup") + }) + + awaitAgent = func(c *codersdk.Client) { + if o.WorkspaceID == uuid.Nil { + require.FailNow(t, "must specify workspace ID for agent in order to wait") + } + coderdtest.AwaitWorkspaceAgents(t, c, o.WorkspaceID) + } + + return agentClient, awaitAgent +} From 3d2788e4ebd85e70981043e760c82d52c606fb56 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 14:44:51 +0000 Subject: [PATCH 3/9] return wrapper instead of func --- agent/agenttest/agent.go | 54 ++++++++++++++++++++++++++++++++-------- 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/agent/agenttest/agent.go b/agent/agenttest/agent.go index fc2086e562de8..c945d5aa50e69 100644 --- a/agent/agenttest/agent.go +++ b/agent/agenttest/agent.go @@ -3,6 +3,7 @@ package agenttest import ( "context" "net/url" + "sync" "testing" "github.com/stretchr/testify/require" @@ -35,6 +36,34 @@ type Options struct { Logger *slog.Logger } +// Agent is a small wrapper around an agent for use in tests. +type Agent struct { + waitOnce sync.Once + agent agent.Agent + agentClient *agentsdk.Client + resources []codersdk.WorkspaceResource + waiter func(*codersdk.Client) []codersdk.WorkspaceResource +} + +// Wait waits for the agent to connect to the workspace and returns the +// resources for the connected workspace. +func (a *Agent) Wait(client *codersdk.Client) []codersdk.WorkspaceResource { + a.waitOnce.Do(func() { + a.resources = a.waiter(client) + }) + 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 +} + // OptFunc is a function that modifies the given options. type OptFunc func(*Options) @@ -57,10 +86,10 @@ func WithWorkspaceID(id uuid.UUID) OptFunc { } // New starts a new agent for use in tests. -// Returns the agent client and a function that will block until the agent is -// connected to the workspace. +// Returns a wrapper around the agent that can be used to wait for the agent to +// connect to the workspace. // Closing the agent is handled by the test cleanup. -func New(t testing.TB, opts ...OptFunc) (agentClient *agentsdk.Client, awaitAgent func(*codersdk.Client)) { +func New(t testing.TB, opts ...OptFunc) *Agent { t.Helper() var o Options @@ -71,7 +100,7 @@ func New(t testing.TB, opts ...OptFunc) (agentClient *agentsdk.Client, awaitAgen if o.URL == nil { require.Fail(t, "must specify URL for agent") } - agentClient = agentsdk.New(o.URL) + agentClient := agentsdk.New(o.URL) if o.AgentToken == "" { o.AgentToken = uuid.NewString() @@ -104,12 +133,15 @@ func New(t testing.TB, opts ...OptFunc) (agentClient *agentsdk.Client, awaitAgen assert.NoError(t, agentCloser.Close(), "failed to close agent during cleanup") }) - awaitAgent = func(c *codersdk.Client) { - if o.WorkspaceID == uuid.Nil { - require.FailNow(t, "must specify workspace ID for agent in order to wait") - } - coderdtest.AwaitWorkspaceAgents(t, c, o.WorkspaceID) + return &Agent{ + agent: agentCloser, + agentClient: agentClient, + waiter: func(c *codersdk.Client) []codersdk.WorkspaceResource { + if o.WorkspaceID == uuid.Nil { + require.FailNow(t, "must specify workspace ID for agent in order to wait") + return nil // unreachable + } + return coderdtest.AwaitWorkspaceAgents(t, c, o.WorkspaceID) + }, } - - return agentClient, awaitAgent } From 5b4cad97e3904f2080293f8cc2dd1c3603538e3a Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 15:09:04 +0000 Subject: [PATCH 4/9] refactor existing unit tests --- cli/configssh_test.go | 20 ++-- cli/gitssh_test.go | 33 ++----- cli/ping_test.go | 19 ++-- cli/portforward_test.go | 30 ++---- cli/speedtest_test.go | 17 ++-- cli/ssh_test.go | 111 +++++++++-------------- cli/vscodessh_test.go | 21 ++--- coderd/activitybump_test.go | 18 ++-- coderd/coderd_test.go | 17 ++-- coderd/coderdtest/coderdtest.go | 2 +- coderd/insights_test.go | 32 +++---- coderd/templates_test.go | 18 ++-- coderd/workspaceagents_test.go | 75 ++++++--------- coderd/workspaceapps/db_test.go | 18 ++-- coderd/workspaces_test.go | 36 +++----- enterprise/coderd/workspaceproxy_test.go | 35 +++---- enterprise/wsproxy/wsproxy_test.go | 33 +++---- scaletest/agentconn/run_test.go | 20 ++-- scaletest/reconnectingpty/run_test.go | 21 ++--- scaletest/workspacetraffic/run_test.go | 34 +++---- 20 files changed, 213 insertions(+), 397 deletions(-) diff --git a/cli/configssh_test.go b/cli/configssh_test.go index 44246da2596e7..af84a5a86c1a8 100644 --- a/cli/configssh_test.go +++ b/cli/configssh_test.go @@ -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" @@ -102,16 +99,11 @@ 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil) require.NoError(t, err) defer agentConn.Close() diff --git a/cli/gitssh_test.go b/cli/gitssh_test.go index 0241b54b4f2af..efb237b301dcb 100644 --- a/cli/gitssh_test.go +++ b/cli/gitssh_test.go @@ -15,18 +15,14 @@ import ( "sync/atomic" "testing" - "cdr.dev/slog" - "cdr.dev/slog/sloggers/slogtest" - - "github.com/coder/coder/v2/agent" - "github.com/coder/coder/v2/codersdk/agentsdk" - "github.com/stretchr/testify/assert" - "github.com/gliderlabs/ssh" "github.com/google/uuid" "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" @@ -64,22 +60,13 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) // start workspace agent - agentLog := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent") - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(agentToken) - agt := agent.New(agent.Options{ - Client: agentClient, - Logger: agentLog, - LogDir: t.TempDir(), - ExchangeToken: func(_ context.Context) (string, error) { - return agentToken, nil - }, - }) - t.Cleanup(func() { - assert.NoError(t, agt.Close(), "failed to close agent in cleanup") - }) - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) - return agentClient, agentToken, pubkey + agt := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ) + agt.Wait(client) + return agt.Client(), agentToken, pubkey } func serveSSHForGitSSH(t *testing.T, handler func(ssh.Session), pubkeys ...gossh.PublicKey) *net.TCPAddr { diff --git a/cli/ping_test.go b/cli/ping_test.go index d054cbf38057a..fd22139786cd0 100644 --- a/cli/ping_test.go +++ b/cli/ping_test.go @@ -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" ) @@ -29,15 +26,11 @@ 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/cli/portforward_test.go b/cli/portforward_test.go index 920980e72fc2e..eda2f67496dcf 100644 --- a/cli/portforward_test.go +++ b/cli/portforward_test.go @@ -14,13 +14,11 @@ 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" @@ -317,24 +315,14 @@ 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + func(o *agenttest.Options) { + o.AgentOptions.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) return workspace } diff --git a/cli/speedtest_test.go b/cli/speedtest_test.go index 5fc9aedbe2e79..fb88a90b75522 100644 --- a/cli/speedtest_test.go +++ b/cli/speedtest_test.go @@ -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" ) @@ -26,14 +24,11 @@ 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 971dc2873ffdc..0a1cfd1a3c519 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -29,13 +29,12 @@ 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/clitest" "github.com/coder/coder/v2/cli/cliui" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "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" @@ -103,15 +102,11 @@ func TestSSH(t *testing.T) { }) pty.ExpectMatch("Waiting") - 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -167,15 +162,11 @@ func TestSSH(t *testing.T) { }) pty.ExpectMatch("Waiting") - 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Ensure the agent is connected. pty.WriteLine("echo hell'o'") @@ -196,14 +187,12 @@ func TestSSH(t *testing.T) { _, _ = tGoContext(t, func(ctx context.Context) { // Run this async so the SSH command has to wait for // the build and agent to connect! - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(agentToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent"), - }) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) <-ctx.Done() - _ = agentCloser.Close() }) clientOutput, clientInput := io.Pipe() @@ -264,14 +253,12 @@ func TestSSH(t *testing.T) { _, _ = tGoContext(t, func(ctx context.Context) { // Run this async so the SSH command has to wait for // the build and agent to connect. - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(agentToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent"), - }) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) <-ctx.Done() - _ = agentCloser.Close() }) clientOutput, clientInput := io.Pipe() @@ -333,13 +320,11 @@ func TestSSH(t *testing.T) { 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"), - }) - defer agentCloser.Close() + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Generate private key. privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) @@ -424,13 +409,11 @@ func TestSSH(t *testing.T) { 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"), - }) - defer agentCloser.Close() + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -475,15 +458,11 @@ func TestSSH(t *testing.T) { pty.ExpectMatch("Waiting") - 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, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -651,16 +630,16 @@ Expire-Date: 0 client, workspace, agentToken := setupWorkspaceForAgent(t, nil) - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(agentToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - EnvironmentVariables: map[string]string{ - "GNUPGHOME": gnupgHomeWorkspace, + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + func(o *agenttest.Options) { + o.AgentOptions.EnvironmentVariables = map[string]string{ + "GNUPGHOME": gnupgHomeWorkspace, + } }, - Logger: slogtest.Make(t, nil).Named("agent"), - }) - defer agentCloser.Close() + ).Wait(client) inv, root := clitest.New(t, "ssh", diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go index 2c1afd6135587..832015b0d2e92 100644 --- a/cli/vscodessh_test.go +++ b/cli/vscodessh_test.go @@ -9,13 +9,9 @@ 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/pty/ptytest" "github.com/coder/coder/v2/testutil" ) @@ -29,16 +25,11 @@ func TestVSCodeSSH(t *testing.T) { user, err := client.User(ctx, codersdk.Me) require.NoError(t, err) - 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() - }() - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) fs := afero.NewMemMapFs() err = afero.WriteFile(fs, "/url", []byte(client.URL.String()), 0o600) diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go index 024803157b784..4bacc1b44ef63 100644 --- a/coderd/activitybump_test.go +++ b/coderd/activitybump_test.go @@ -9,14 +9,13 @@ import ( "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/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/schedule" "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/testutil" ) @@ -87,16 +86,11 @@ func TestWorkspaceActivityBump(t *testing.T) { require.NoError(t, err) } - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(agentToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent"), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Sanity-check that deadline is near. workspace, err := client.Workspace(ctx, workspace.ID) diff --git a/coderd/coderd_test.go b/coderd/coderd_test.go index 6edf4657cc903..3b9edae5f4843 100644 --- a/coderd/coderd_test.go +++ b/coderd/coderd_test.go @@ -22,12 +22,11 @@ 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/buildinfo" "github.com/coder/coder/v2/coderd" "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/tailnet" "github.com/coder/coder/v2/testutil" @@ -189,15 +188,11 @@ func TestDERPForceWebSockets(t *testing.T) { 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").Leveled(slog.LevelDebug), - }) - defer func() { - _ = agentCloser.Close() - }() + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 5a00a679d808e..142c1341cab19 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -798,7 +798,7 @@ func AwaitWorkspaceBuildJob(t *testing.T, client *codersdk.Client, build uuid.UU // AwaitWorkspaceAgents waits for all resources with agents to be connected. If // specific agents are provided, it will wait for those agents to be connected // but will not fail if other agents are not connected. -func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource { +func AwaitWorkspaceAgents(t testing.TB, client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource { t.Helper() agentNamesMap := make(map[string]struct{}, len(agentNames)) diff --git a/coderd/insights_test.go b/coderd/insights_test.go index 8be3f4b5e6076..7d94d3ea49867 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -19,7 +19,7 @@ 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/coderd/batchstats" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" @@ -58,16 +58,11 @@ func TestDeploymentInsights(t *testing.T) { 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{ - Logger: slogtest.Make(t, nil), - Client: agentClient, - }) - defer func() { - _ = agentCloser.Close() - }() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -145,16 +140,11 @@ func TestUserLatencyInsights(t *testing.T) { coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) // Start an agent so that we can generate stats. - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Logger: logger.Named("agent"), - Client: agentClient, - }) - defer func() { - _ = agentCloser.Close() - }() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Start must be at the beginning of the day, initialize it early in case // the day changes so that we get the relevant stats faster. diff --git a/coderd/templates_test.go b/coderd/templates_test.go index 528485ac2faf5..540299a55c5ce 100644 --- a/coderd/templates_test.go +++ b/coderd/templates_test.go @@ -12,7 +12,7 @@ import ( "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/coderd/audit" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" @@ -20,7 +20,6 @@ import ( "github.com/coder/coder/v2/coderd/schedule" "github.com/coder/coder/v2/coderd/util/ptr" "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/testutil" ) @@ -1218,16 +1217,11 @@ func TestTemplateMetrics(t *testing.T) { 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{ - Logger: slogtest.Make(t, nil), - Client: agentClient, - }) - defer func() { - _ = agentCloser.Close() - }() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 1a0e5e2cd7b54..279c290abbe41 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -20,7 +20,7 @@ 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/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtime" @@ -484,20 +484,15 @@ func TestWorkspaceAgentListen(t *testing.T) { 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").Leveled(slog.LevelDebug), - }) - defer func() { - _ = agentCloser.Close() - }() + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) conn, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID, nil) require.NoError(t, err) defer func() { @@ -585,14 +580,11 @@ func TestWorkspaceAgentTailnet(t *testing.T) { coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) daemonCloser.Close() - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug), - }) - defer agentCloser.Close() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancelFunc := context.WithCancel(context.Background()) defer cancelFunc() @@ -646,12 +638,11 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) { require.NoError(t, err) require.True(t, manifest.DisableDirectConnections) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug), - }) - defer agentCloser.Close() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) agentID := resources[0].Agents[0].ID // Verify that the connection data has no STUN ports and @@ -729,17 +720,11 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) { 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").Leveled(slog.LevelDebug), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) - + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) return client, uint16(coderdPort), resources[0].Agents[0].ID } @@ -1457,17 +1442,15 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) { coderdtest.AwaitTemplateVersionJob(t, client, 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(agentToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: logger.Named("agent"), - }) - defer func() { - _ = agentCloser.Close() - }() - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + + agt := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentToken), + agenttest.WithWorkspaceID(workspace.ID), + ) + resources := agt.Wait(client) agentID := resources[0].Agents[0].ID + agentCloser := agt.Agent() // Connect from a client. ctx := testutil.Context(t, testutil.WaitLong) diff --git a/coderd/workspaceapps/db_test.go b/coderd/workspaceapps/db_test.go index 3be987e28ca23..d14bc4e34ac73 100644 --- a/coderd/workspaceapps/db_test.go +++ b/coderd/workspaceapps/db_test.go @@ -17,14 +17,12 @@ 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/coderd/coderdtest" "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/coderd/workspaceapps" "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/testutil" @@ -173,15 +171,11 @@ func Test_ResolveRequest(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, firstUser.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(agentAuthToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent"), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(agentAuthToken), + agenttest.WithWorkspaceID(workspace.ID), + ) resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID, agentName) agentID := uuid.Nil diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index 83884f862489c..45922629d2769 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -19,7 +19,7 @@ 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/coderd/audit" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" @@ -33,7 +33,6 @@ import ( "github.com/coder/coder/v2/coderd/schedule/cron" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" - "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/provisioner/echo" "github.com/coder/coder/v2/provisionersdk/proto" @@ -1348,17 +1347,11 @@ func TestWorkspaceFilterManual(t *testing.T) { 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").Leveled(slog.LevelDebug), - }) - defer func() { - _ = agentCloser.Close() - }() - - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -2268,21 +2261,18 @@ func TestWorkspaceWatcher(t *testing.T) { wait("agent timeout after create", nil) wait("agent timeout after start", nil) - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: logger.Named("agent"), - }) - defer func() { - _ = agentCloser.Close() - }() + agt := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ) + agt.Wait(client) wait("agent connected/ready", func(w codersdk.Workspace) bool { return w.LatestBuild.Resources[0].Agents[0].Status == codersdk.WorkspaceAgentConnected && w.LatestBuild.Resources[0].Agents[0].LifecycleState == codersdk.WorkspaceAgentLifecycleReady }) - agentCloser.Close() + agt.Agent().Close() wait("agent disconnected", func(w codersdk.Workspace) bool { return w.LatestBuild.Resources[0].Agents[0].Status == codersdk.WorkspaceAgentDisconnected }) diff --git a/enterprise/coderd/workspaceproxy_test.go b/enterprise/coderd/workspaceproxy_test.go index 745a1af2c11bc..8d9d3f05ff695 100644 --- a/enterprise/coderd/workspaceproxy_test.go +++ b/enterprise/coderd/workspaceproxy_test.go @@ -14,16 +14,13 @@ 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/buildinfo" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/workspaceapps" "github.com/coder/coder/v2/codersdk" - "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" "github.com/coder/coder/v2/enterprise/coderd/license" "github.com/coder/coder/v2/enterprise/wsproxy/wsproxysdk" @@ -677,15 +674,12 @@ func TestIssueSignedAppToken(t *testing.T) { workspace.LatestBuild = build // Connect an agent to the workspace - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) + + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) @@ -787,16 +781,11 @@ func TestReconnectingPTYSignedToken(t *testing.T) { // Connect an agent to the workspace agentID := build.Resources[0].Agents[0].ID - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) proxyURL, err := url.Parse(fmt.Sprintf("https://%s.com", namesgenerator.GetRandomName(1))) require.NoError(t, err) diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go index 3b2caf38f1254..68676bd0b6185 100644 --- a/enterprise/wsproxy/wsproxy_test.go +++ b/enterprise/wsproxy/wsproxy_test.go @@ -15,7 +15,7 @@ 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/clibase" "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/coderdtest" @@ -23,7 +23,6 @@ import ( "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/coderd/workspaceapps/apptest" "github.com/coder/coder/v2/codersdk" - "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" "github.com/coder/coder/v2/enterprise/coderd/license" "github.com/coder/coder/v2/provisioner/echo" @@ -184,16 +183,11 @@ resourceLoop: require.NotEqual(t, uuid.Nil, agentID) // Connect an agent to the workspace - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) t.Run("ReturnedInDERPMap", func(t *testing.T) { t.Parallel() @@ -417,16 +411,11 @@ resourceLoop: require.NotEqual(t, uuid.Nil, agentID) // Connect an agent to the workspace - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) // Connect to the workspace agent. ctx := testutil.Context(t, testutil.WaitLong) diff --git a/scaletest/agentconn/run_test.go b/scaletest/agentconn/run_test.go index 4d3ffb8d0da2d..5aa3f3983a389 100644 --- a/scaletest/agentconn/run_test.go +++ b/scaletest/agentconn/run_test.go @@ -13,12 +13,10 @@ import ( "github.com/google/uuid" "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/coderd/coderdtest" "github.com/coder/coder/v2/coderd/httpapi" "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/scaletest/agentconn" @@ -258,17 +256,11 @@ func setupRunnerTest(t *testing.T) (client *codersdk.Client, agentID uuid.UUID) 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"), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) return client, resources[0].Agents[0].ID } diff --git a/scaletest/reconnectingpty/run_test.go b/scaletest/reconnectingpty/run_test.go index 81de3dcfb9da8..4dedacbc1c73c 100644 --- a/scaletest/reconnectingpty/run_test.go +++ b/scaletest/reconnectingpty/run_test.go @@ -9,13 +9,10 @@ import ( "github.com/google/uuid" "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/coderd/coderdtest" "github.com/coder/coder/v2/coderd/httpapi" "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/scaletest/reconnectingpty" @@ -279,17 +276,11 @@ func setupRunnerTest(t *testing.T) (client *codersdk.Client, agentID uuid.UUID) 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, &slogtest.Options{IgnoreErrors: true}).Named("agent").Leveled(slog.LevelDebug), - }) - t.Cleanup(func() { - _ = agentCloser.Close() - }) - - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(workspace.ID), + ).Wait(client) require.Eventually(t, func() bool { t.Log("agent id", resources[0].Agents[0].ID) return (*api.TailnetCoordinator.Load()).Node(resources[0].Agents[0].ID) != nil diff --git a/scaletest/workspacetraffic/run_test.go b/scaletest/workspacetraffic/run_test.go index 961263e972e1c..0f281cd261877 100644 --- a/scaletest/workspacetraffic/run_test.go +++ b/scaletest/workspacetraffic/run_test.go @@ -10,10 +10,9 @@ import ( "golang.org/x/exp/slices" - "github.com/coder/coder/v2/agent" + "github.com/coder/coder/v2/agent/agenttest" "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/scaletest/workspacetraffic" @@ -73,20 +72,15 @@ func TestRun(t *testing.T) { ) // We also need a running agent to run this test. - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - }) - + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(ws.ID), + ).Wait(client) ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) - t.Cleanup(func() { - _ = agentCloser.Close() - }) // Make sure the agent is connected before we go any further. - resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID) var agentID uuid.UUID for _, res := range resources { for _, agt := range res.Agents { @@ -199,20 +193,16 @@ func TestRun(t *testing.T) { ) // We also need a running agent to run this test. - agentClient := agentsdk.New(client.URL) - agentClient.SetSessionToken(authToken) - agentCloser := agent.New(agent.Options{ - Client: agentClient, - }) + resources := agenttest.New(t, + agenttest.WithURL(client.URL), + agenttest.WithAgentToken(authToken), + agenttest.WithWorkspaceID(ws.ID), + ).Wait(client) ctx, cancel := context.WithCancel(context.Background()) - t.Cleanup(func() { - cancel() - _ = agentCloser.Close() - }) + t.Cleanup(cancel) // Make sure the agent is connected before we go any further. - resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID) var agentID uuid.UUID for _, res := range resources { for _, agt := range res.Agents { From 81a3a78941159132dd59cabdf2729562df1c888e Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 15:22:01 +0000 Subject: [PATCH 5/9] ignore duplicated code in unit tests --- scaletest/workspacetraffic/run_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scaletest/workspacetraffic/run_test.go b/scaletest/workspacetraffic/run_test.go index 0f281cd261877..2064a3abccefb 100644 --- a/scaletest/workspacetraffic/run_test.go +++ b/scaletest/workspacetraffic/run_test.go @@ -32,6 +32,7 @@ func TestRun(t *testing.T) { t.Skip("Race detector enabled, skipping time-sensitive test.") } + //nolint:dupl t.Run("PTY", func(t *testing.T) { t.Parallel() // We need to stand up an in-memory coderd and run a fake workspace. @@ -153,6 +154,7 @@ func TestRun(t *testing.T) { assert.Zero(t, writeMetrics.Errors()) }) + //nolint:dupl t.Run("SSH", func(t *testing.T) { t.Parallel() // We need to stand up an in-memory coderd and run a fake workspace. From 3316b99d1b884a9fc422b3d6a84fd9b8b4141776 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 20:18:14 +0000 Subject: [PATCH 6/9] address code review comments --- agent/agenttest/agent.go | 109 +++++++---------------- cli/configssh_test.go | 6 +- cli/gitssh_test.go | 8 +- cli/ping_test.go | 6 +- cli/portforward_test.go | 12 ++- cli/speedtest_test.go | 6 +- cli/ssh_test.go | 59 +++--------- cli/vscodessh_test.go | 6 +- coderd/activitybump_test.go | 6 +- coderd/coderd_test.go | 6 +- coderd/insights_test.go | 13 +-- coderd/templates_test.go | 6 +- coderd/workspaceagents_test.go | 32 ++----- coderd/workspaceapps/db_test.go | 7 +- coderd/workspaces_test.go | 14 +-- enterprise/coderd/workspaceproxy_test.go | 12 +-- enterprise/wsproxy/wsproxy_test.go | 12 +-- scaletest/agentconn/run_test.go | 6 +- scaletest/reconnectingpty/run_test.go | 6 +- scaletest/workspacetraffic/run_test.go | 12 +-- 20 files changed, 78 insertions(+), 266 deletions(-) diff --git a/agent/agenttest/agent.go b/agent/agenttest/agent.go index c945d5aa50e69..198dfc3bca933 100644 --- a/agent/agenttest/agent.go +++ b/agent/agenttest/agent.go @@ -6,8 +6,6 @@ import ( "sync" "testing" - "github.com/stretchr/testify/require" - "cdr.dev/slog" "cdr.dev/slog/sloggers/slogtest" @@ -20,36 +18,22 @@ import ( "github.com/coder/coder/v2/codersdk/agentsdk" ) -// Options are options for creating a new test agent. -type Options struct { - // AgentOptions are the options to use for the agent. - AgentOptions agent.Options - - // AgentToken is the token to use for the agent. - AgentToken string - // URL is the URL to which the agent should connect. - URL *url.URL - // WorkspaceID is the ID of the workspace to which the agent should connect. - WorkspaceID uuid.UUID - // Logger is the logger to use for the agent. - // Defaults to a new test logger if not specified. - Logger *slog.Logger -} - // Agent is a small wrapper around an agent for use in tests. type Agent struct { - waitOnce sync.Once agent agent.Agent agentClient *agentsdk.Client resources []codersdk.WorkspaceResource - waiter func(*codersdk.Client) []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. -func (a *Agent) Wait(client *codersdk.Client) []codersdk.WorkspaceResource { +// 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) + a.resources = a.waiter(client, workspaceID, agentNames...) }) return a.resources } @@ -64,84 +48,51 @@ func (a *Agent) Agent() agent.Agent { return a.agent } -// OptFunc is a function that modifies the given options. -type OptFunc func(*Options) - -func WithAgentToken(token string) OptFunc { - return func(o *Options) { - o.AgentToken = token - } -} - -func WithURL(u *url.URL) OptFunc { - return func(o *Options) { - o.URL = u - } -} - -func WithWorkspaceID(id uuid.UUID) OptFunc { - return func(o *Options) { - o.WorkspaceID = id - } -} - // New starts a new agent for use in tests. -// Returns a wrapper around the agent that can be used to wait for the agent to -// connect to the workspace. +// 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, opts ...OptFunc) *Agent { +func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent.Options)) *Agent { t.Helper() - var o Options + 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.URL == nil { - require.Fail(t, "must specify URL for agent") - } - agentClient := agentsdk.New(o.URL) - - if o.AgentToken == "" { - o.AgentToken = uuid.NewString() + if o.Client == nil { + agentClient := agentsdk.New(coderURL) + agentClient.SetSessionToken(agentToken) + o.Client = agentClient } - agentClient.SetSessionToken(o.AgentToken) - if o.AgentOptions.Client == nil { - o.AgentOptions.Client = agentClient - } - - if o.AgentOptions.ExchangeToken == nil { - o.AgentOptions.ExchangeToken = func(_ context.Context) (string, error) { - return o.AgentToken, nil + if o.ExchangeToken == nil { + o.ExchangeToken = func(_ context.Context) (string, error) { + return agentToken, nil } } - if o.AgentOptions.LogDir == "" { - o.AgentOptions.LogDir = t.TempDir() + if o.LogDir == "" { + o.LogDir = t.TempDir() } - if o.Logger == nil { - log := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("agent") - o.Logger = &log - } - - o.AgentOptions.Logger = *o.Logger - - agentCloser := agent.New(o.AgentOptions) + agentCloser := agent.New(o) t.Cleanup(func() { assert.NoError(t, agentCloser.Close(), "failed to close agent during cleanup") }) return &Agent{ agent: agentCloser, - agentClient: agentClient, - waiter: func(c *codersdk.Client) []codersdk.WorkspaceResource { - if o.WorkspaceID == uuid.Nil { - require.FailNow(t, "must specify workspace ID for agent in order to wait") - return nil // unreachable - } - return coderdtest.AwaitWorkspaceAgents(t, c, o.WorkspaceID) + 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...) }, } } diff --git a/cli/configssh_test.go b/cli/configssh_test.go index af84a5a86c1a8..41515fcad5f9b 100644 --- a/cli/configssh_test.go +++ b/cli/configssh_test.go @@ -99,11 +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) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + 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() diff --git a/cli/gitssh_test.go b/cli/gitssh_test.go index efb237b301dcb..ce41e853712b3 100644 --- a/cli/gitssh_test.go +++ b/cli/gitssh_test.go @@ -60,12 +60,8 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) // start workspace agent - agt := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ) - agt.Wait(client) + agt := agenttest.New(t, client.URL, agentToken) + agt.Wait(client, workspace.ID) return agt.Client(), agentToken, pubkey } diff --git a/cli/ping_test.go b/cli/ping_test.go index fd22139786cd0..24d64d660b362 100644 --- a/cli/ping_test.go +++ b/cli/ping_test.go @@ -26,11 +26,7 @@ func TestPing(t *testing.T) { inv.Stderr = pty.Output() inv.Stdout = pty.Output() - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/cli/portforward_test.go b/cli/portforward_test.go index eda2f67496dcf..df31c2b32b4d0 100644 --- a/cli/portforward_test.go +++ b/cli/portforward_test.go @@ -14,6 +14,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/coder/coder/v2/agent" "github.com/coder/coder/v2/agent/agenttest" "github.com/coder/coder/v2/cli/clitest" @@ -315,14 +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) - agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - func(o *agenttest.Options) { - o.AgentOptions.SSHMaxTimeout = 60 * time.Second + agenttest.New(t, client.URL, agentToken, + func(o *agent.Options) { + o.SSHMaxTimeout = 60 * time.Second }, - ).Wait(client) + ).Wait(client, workspace.ID) return workspace } diff --git a/cli/speedtest_test.go b/cli/speedtest_test.go index fb88a90b75522..d28eb5bed8cb3 100644 --- a/cli/speedtest_test.go +++ b/cli/speedtest_test.go @@ -24,11 +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) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 0a1cfd1a3c519..388eb8542ca20 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -29,6 +29,7 @@ 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/clitest" "github.com/coder/coder/v2/cli/cliui" @@ -102,11 +103,7 @@ func TestSSH(t *testing.T) { }) pty.ExpectMatch("Waiting") - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -162,11 +159,7 @@ func TestSSH(t *testing.T) { }) pty.ExpectMatch("Waiting") - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) // Ensure the agent is connected. pty.WriteLine("echo hell'o'") @@ -187,11 +180,7 @@ func TestSSH(t *testing.T) { _, _ = 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, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) <-ctx.Done() }) @@ -253,11 +242,7 @@ func TestSSH(t *testing.T) { _, _ = 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, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) <-ctx.Done() }) @@ -320,11 +305,7 @@ func TestSSH(t *testing.T) { client, workspace, agentToken := setupWorkspaceForAgent(t, nil) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) // Generate private key. privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) @@ -409,11 +390,7 @@ func TestSSH(t *testing.T) { client, workspace, agentToken := setupWorkspaceForAgent(t, nil) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -458,11 +435,7 @@ func TestSSH(t *testing.T) { pty.ExpectMatch("Waiting") - agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -630,16 +603,12 @@ Expire-Date: 0 client, workspace, agentToken := setupWorkspaceForAgent(t, nil) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - func(o *agenttest.Options) { - o.AgentOptions.EnvironmentVariables = map[string]string{ - "GNUPGHOME": gnupgHomeWorkspace, - } - }, - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) { + o.EnvironmentVariables = map[string]string{ + "GNUPGHOME": gnupgHomeWorkspace, + } + }, + ).Wait(client, workspace.ID) inv, root := clitest.New(t, "ssh", diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go index 832015b0d2e92..73d6a5e05017e 100644 --- a/cli/vscodessh_test.go +++ b/cli/vscodessh_test.go @@ -25,11 +25,7 @@ func TestVSCodeSSH(t *testing.T) { user, err := client.User(ctx, codersdk.Me) require.NoError(t, err) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) fs := afero.NewMemMapFs() err = afero.WriteFile(fs, "/url", []byte(client.URL.String()), 0o600) diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go index 4bacc1b44ef63..fc0cf92ecb647 100644 --- a/coderd/activitybump_test.go +++ b/coderd/activitybump_test.go @@ -86,11 +86,7 @@ func TestWorkspaceActivityBump(t *testing.T) { require.NoError(t, err) } - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) // Sanity-check that deadline is near. workspace, err := client.Workspace(ctx, workspace.ID) diff --git a/coderd/coderd_test.go b/coderd/coderd_test.go index 3b9edae5f4843..7edd769fa0abb 100644 --- a/coderd/coderd_test.go +++ b/coderd/coderd_test.go @@ -188,11 +188,7 @@ func TestDERPForceWebSockets(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/coderd/insights_test.go b/coderd/insights_test.go index 7d94d3ea49867..e2ca5c8a5c5ad 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -58,12 +58,7 @@ func TestDeploymentInsights(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) - + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -140,11 +135,7 @@ func TestUserLatencyInsights(t *testing.T) { coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) // Start an agent so that we can generate stats. - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) // Start must be at the beginning of the day, initialize it early in case // the day changes so that we get the relevant stats faster. diff --git a/coderd/templates_test.go b/coderd/templates_test.go index 540299a55c5ce..7c8c72c6860ed 100644 --- a/coderd/templates_test.go +++ b/coderd/templates_test.go @@ -1217,11 +1217,7 @@ func TestTemplateMetrics(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 279c290abbe41..893974790660c 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -484,11 +484,7 @@ func TestWorkspaceAgentListen(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -580,11 +576,7 @@ func TestWorkspaceAgentTailnet(t *testing.T) { coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) daemonCloser.Close() - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) ctx, cancelFunc := context.WithCancel(context.Background()) defer cancelFunc() @@ -638,11 +630,7 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) { require.NoError(t, err) require.True(t, manifest.DisableDirectConnections) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) agentID := resources[0].Agents[0].ID // Verify that the connection data has no STUN ports and @@ -720,11 +708,7 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) return client, uint16(coderdPort), resources[0].Agents[0].ID } @@ -1443,12 +1427,8 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - agt := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentToken), - agenttest.WithWorkspaceID(workspace.ID), - ) - resources := agt.Wait(client) + agt := agenttest.New(t, client.URL, agentToken) + resources := agt.Wait(client, workspace.ID) agentID := resources[0].Agents[0].ID agentCloser := agt.Agent() diff --git a/coderd/workspaceapps/db_test.go b/coderd/workspaceapps/db_test.go index d14bc4e34ac73..06e0bfa8dead7 100644 --- a/coderd/workspaceapps/db_test.go +++ b/coderd/workspaceapps/db_test.go @@ -171,12 +171,7 @@ func Test_ResolveRequest(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, firstUser.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(agentAuthToken), - agenttest.WithWorkspaceID(workspace.ID), - ) - resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID, agentName) + resources := agenttest.New(t, client.URL, agentAuthToken).Wait(client, workspace.ID, agentName) agentID := uuid.Nil for _, resource := range resources { diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index 45922629d2769..90166e71d5e6f 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -1347,11 +1347,7 @@ func TestWorkspaceFilterManual(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -2261,12 +2257,8 @@ func TestWorkspaceWatcher(t *testing.T) { wait("agent timeout after create", nil) wait("agent timeout after start", nil) - agt := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ) - agt.Wait(client) + agt := agenttest.New(t, client.URL, authToken) + agt.Wait(client, workspace.ID) wait("agent connected/ready", func(w codersdk.Workspace) bool { return w.LatestBuild.Resources[0].Agents[0].Status == codersdk.WorkspaceAgentConnected && diff --git a/enterprise/coderd/workspaceproxy_test.go b/enterprise/coderd/workspaceproxy_test.go index 8d9d3f05ff695..40cd2592ba743 100644 --- a/enterprise/coderd/workspaceproxy_test.go +++ b/enterprise/coderd/workspaceproxy_test.go @@ -675,11 +675,7 @@ func TestIssueSignedAppToken(t *testing.T) { // Connect an agent to the workspace - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) @@ -781,11 +777,7 @@ func TestReconnectingPTYSignedToken(t *testing.T) { // Connect an agent to the workspace agentID := build.Resources[0].Agents[0].ID - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) proxyURL, err := url.Parse(fmt.Sprintf("https://%s.com", namesgenerator.GetRandomName(1))) require.NoError(t, err) diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go index 68676bd0b6185..18b1017df2036 100644 --- a/enterprise/wsproxy/wsproxy_test.go +++ b/enterprise/wsproxy/wsproxy_test.go @@ -183,11 +183,7 @@ resourceLoop: require.NotEqual(t, uuid.Nil, agentID) // Connect an agent to the workspace - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) t.Run("ReturnedInDERPMap", func(t *testing.T) { t.Parallel() @@ -411,11 +407,7 @@ resourceLoop: require.NotEqual(t, uuid.Nil, agentID) // Connect an agent to the workspace - _ = agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) // Connect to the workspace agent. ctx := testutil.Context(t, testutil.WaitLong) diff --git a/scaletest/agentconn/run_test.go b/scaletest/agentconn/run_test.go index 5aa3f3983a389..36dae72ef6e82 100644 --- a/scaletest/agentconn/run_test.go +++ b/scaletest/agentconn/run_test.go @@ -256,11 +256,7 @@ func setupRunnerTest(t *testing.T) (client *codersdk.Client, agentID uuid.UUID) workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) return client, resources[0].Agents[0].ID } diff --git a/scaletest/reconnectingpty/run_test.go b/scaletest/reconnectingpty/run_test.go index 4dedacbc1c73c..b85df2437a1a9 100644 --- a/scaletest/reconnectingpty/run_test.go +++ b/scaletest/reconnectingpty/run_test.go @@ -276,11 +276,7 @@ func setupRunnerTest(t *testing.T) (client *codersdk.Client, agentID uuid.UUID) workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(workspace.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) require.Eventually(t, func() bool { t.Log("agent id", resources[0].Agents[0].ID) return (*api.TailnetCoordinator.Load()).Node(resources[0].Agents[0].ID) != nil diff --git a/scaletest/workspacetraffic/run_test.go b/scaletest/workspacetraffic/run_test.go index 2064a3abccefb..604b606806569 100644 --- a/scaletest/workspacetraffic/run_test.go +++ b/scaletest/workspacetraffic/run_test.go @@ -73,11 +73,7 @@ func TestRun(t *testing.T) { ) // We also need a running agent to run this test. - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(ws.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, ws.ID) ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) @@ -195,11 +191,7 @@ func TestRun(t *testing.T) { ) // We also need a running agent to run this test. - resources := agenttest.New(t, - agenttest.WithURL(client.URL), - agenttest.WithAgentToken(authToken), - agenttest.WithWorkspaceID(ws.ID), - ).Wait(client) + resources := agenttest.New(t, client.URL, authToken).Wait(client, ws.ID) ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) From c2cba9dd87e8263cbb654aabcfc1a8950d0328ee Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 21 Sep 2023 20:18:25 +0000 Subject: [PATCH 7/9] fixup! address code review comments --- agent/agenttest/agent.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/agent/agenttest/agent.go b/agent/agenttest/agent.go index 198dfc3bca933..e382b50c1c6c1 100644 --- a/agent/agenttest/agent.go +++ b/agent/agenttest/agent.go @@ -6,12 +6,11 @@ import ( "sync" "testing" - "cdr.dev/slog" - "cdr.dev/slog/sloggers/slogtest" - "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" From d55c79790e9b21dfd490ebe1e3db77ce114d0488 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Mon, 25 Sep 2023 12:38:04 +0100 Subject: [PATCH 8/9] address PR comments --- agent/agenttest/agent.go | 56 ++++-------------------- cli/configssh_test.go | 3 +- cli/gitssh_test.go | 14 +++--- cli/ping_test.go | 4 +- cli/portforward_test.go | 6 +-- cli/speedtest_test.go | 4 +- cli/ssh_test.go | 25 +++++++---- cli/vscodessh_test.go | 4 +- coderd/activitybump_test.go | 3 +- coderd/coderd_test.go | 3 +- coderd/insights_test.go | 6 ++- coderd/templates_test.go | 3 +- coderd/workspaceagents_test.go | 17 ++++--- coderd/workspaceapps/db_test.go | 3 +- coderd/workspaces_test.go | 7 +-- enterprise/coderd/workspaceproxy_test.go | 9 ++-- enterprise/wsproxy/wsproxy_test.go | 6 ++- scaletest/agentconn/run_test.go | 3 +- scaletest/reconnectingpty/run_test.go | 3 +- scaletest/workspacetraffic/run_test.go | 6 ++- 20 files changed, 89 insertions(+), 96 deletions(-) diff --git a/agent/agenttest/agent.go b/agent/agenttest/agent.go index e382b50c1c6c1..77b7c6e368822 100644 --- a/agent/agenttest/agent.go +++ b/agent/agenttest/agent.go @@ -3,59 +3,24 @@ 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 { +// Returns the agent. Closing the agent is handled by the test cleanup. +// It is the responsibility of the caller to call coderdtest.AwaitWorkspaceAgents +// to ensure agent is connected. +func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent.Options)) agent.Agent { t.Helper() var o agent.Options @@ -69,6 +34,7 @@ func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent if o.Client == nil { agentClient := agentsdk.New(coderURL) agentClient.SetSessionToken(agentToken) + agentClient.SDK.SetLogger(log) o.Client = agentClient } @@ -82,16 +48,10 @@ func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent o.LogDir = t.TempDir() } - agentCloser := agent.New(o) + agt := agent.New(o) t.Cleanup(func() { - assert.NoError(t, agentCloser.Close(), "failed to close agent during cleanup") + assert.NoError(t, agt.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...) - }, - } + return agt } diff --git a/cli/configssh_test.go b/cli/configssh_test.go index 41515fcad5f9b..141609f25bf53 100644 --- a/cli/configssh_test.go +++ b/cli/configssh_test.go @@ -99,7 +99,8 @@ 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) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil) require.NoError(t, err) defer agentConn.Close() diff --git a/cli/gitssh_test.go b/cli/gitssh_test.go index ce41e853712b3..103ce34cf723b 100644 --- a/cli/gitssh_test.go +++ b/cli/gitssh_test.go @@ -20,12 +20,12 @@ import ( "github.com/stretchr/testify/require" gossh "golang.org/x/crypto/ssh" + "github.com/coder/coder/v2/agent" "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" + "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" @@ -60,9 +60,13 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) // start workspace agent - agt := agenttest.New(t, client.URL, agentToken) - agt.Wait(client, workspace.ID) - return agt.Client(), agentToken, pubkey + agentClient := agentsdk.New(client.URL) + agentClient.SetSessionToken(agentToken) + _ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) { + o.Client = agentClient + }) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + return agentClient, agentToken, pubkey } func serveSSHForGitSSH(t *testing.T, handler func(ssh.Session), pubkeys ...gossh.PublicKey) *net.TCPAddr { diff --git a/cli/ping_test.go b/cli/ping_test.go index 24d64d660b362..f2bd4b5ff88a1 100644 --- a/cli/ping_test.go +++ b/cli/ping_test.go @@ -8,6 +8,7 @@ import ( "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/pty/ptytest" "github.com/coder/coder/v2/testutil" ) @@ -26,7 +27,8 @@ func TestPing(t *testing.T) { inv.Stderr = pty.Output() inv.Stdout = pty.Output() - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/cli/portforward_test.go b/cli/portforward_test.go index df31c2b32b4d0..d2d497485c403 100644 --- a/cli/portforward_test.go +++ b/cli/portforward_test.go @@ -16,7 +16,6 @@ import ( "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" @@ -316,11 +315,12 @@ 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) - agenttest.New(t, client.URL, agentToken, + _ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) { o.SSHMaxTimeout = 60 * time.Second }, - ).Wait(client, workspace.ID) + ) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) return workspace } diff --git a/cli/speedtest_test.go b/cli/speedtest_test.go index d28eb5bed8cb3..f16b769cc85e7 100644 --- a/cli/speedtest_test.go +++ b/cli/speedtest_test.go @@ -12,6 +12,7 @@ import ( "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/pty/ptytest" "github.com/coder/coder/v2/testutil" @@ -24,7 +25,8 @@ 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) - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 388eb8542ca20..2a371d5d0544f 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -103,7 +103,8 @@ func TestSSH(t *testing.T) { }) pty.ExpectMatch("Waiting") - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -159,7 +160,8 @@ func TestSSH(t *testing.T) { }) pty.ExpectMatch("Waiting") - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Ensure the agent is connected. pty.WriteLine("echo hell'o'") @@ -180,7 +182,8 @@ func TestSSH(t *testing.T) { _, _ = 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).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) <-ctx.Done() }) @@ -242,7 +245,8 @@ func TestSSH(t *testing.T) { _, _ = 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).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) <-ctx.Done() }) @@ -305,7 +309,8 @@ func TestSSH(t *testing.T) { client, workspace, agentToken := setupWorkspaceForAgent(t, nil) - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Generate private key. privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) @@ -390,7 +395,8 @@ func TestSSH(t *testing.T) { client, workspace, agentToken := setupWorkspaceForAgent(t, nil) - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -435,7 +441,8 @@ func TestSSH(t *testing.T) { pty.ExpectMatch("Waiting") - agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -607,8 +614,8 @@ Expire-Date: 0 o.EnvironmentVariables = map[string]string{ "GNUPGHOME": gnupgHomeWorkspace, } - }, - ).Wait(client, workspace.ID) + }) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) inv, root := clitest.New(t, "ssh", diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go index 73d6a5e05017e..dc3e65f5bdb9c 100644 --- a/cli/vscodessh_test.go +++ b/cli/vscodessh_test.go @@ -11,6 +11,7 @@ import ( "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/pty/ptytest" "github.com/coder/coder/v2/testutil" @@ -25,7 +26,8 @@ func TestVSCodeSSH(t *testing.T) { user, err := client.User(ctx, codersdk.Me) require.NoError(t, err) - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) fs := afero.NewMemMapFs() err = afero.WriteFile(fs, "/url", []byte(client.URL.String()), 0o600) diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go index 91066b7d9a4dc..53792434bb7bc 100644 --- a/coderd/activitybump_test.go +++ b/coderd/activitybump_test.go @@ -85,7 +85,8 @@ func TestWorkspaceActivityBump(t *testing.T) { require.NoError(t, err) } - _ = agenttest.New(t, client.URL, agentToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Sanity-check that deadline is near. workspace, err := client.Workspace(ctx, workspace.ID) diff --git a/coderd/coderd_test.go b/coderd/coderd_test.go index 7edd769fa0abb..f767e3a7d3571 100644 --- a/coderd/coderd_test.go +++ b/coderd/coderd_test.go @@ -188,7 +188,8 @@ func TestDERPForceWebSockets(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/coderd/insights_test.go b/coderd/insights_test.go index e2ca5c8a5c5ad..d9b4832ec8e43 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -58,7 +58,8 @@ func TestDeploymentInsights(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -135,7 +136,8 @@ func TestUserLatencyInsights(t *testing.T) { coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) // Start an agent so that we can generate stats. - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Start must be at the beginning of the day, initialize it early in case // the day changes so that we get the relevant stats faster. diff --git a/coderd/templates_test.go b/coderd/templates_test.go index 7c8c72c6860ed..3e5ad54e733f6 100644 --- a/coderd/templates_test.go +++ b/coderd/templates_test.go @@ -1217,7 +1217,8 @@ func TestTemplateMetrics(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 893974790660c..b580efb71b3cd 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -484,7 +484,8 @@ func TestWorkspaceAgentListen(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -576,7 +577,8 @@ func TestWorkspaceAgentTailnet(t *testing.T) { coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) daemonCloser.Close() - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancelFunc := context.WithCancel(context.Background()) defer cancelFunc() @@ -630,7 +632,8 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) { require.NoError(t, err) require.True(t, manifest.DisableDirectConnections) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) agentID := resources[0].Agents[0].ID // Verify that the connection data has no STUN ports and @@ -708,7 +711,8 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) return client, uint16(coderdPort), resources[0].Agents[0].ID } @@ -1427,10 +1431,9 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - agt := agenttest.New(t, client.URL, agentToken) - resources := agt.Wait(client, workspace.ID) + agentCloser := agenttest.New(t, client.URL, agentToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) agentID := resources[0].Agents[0].ID - agentCloser := agt.Agent() // Connect from a client. ctx := testutil.Context(t, testutil.WaitLong) diff --git a/coderd/workspaceapps/db_test.go b/coderd/workspaceapps/db_test.go index 06e0bfa8dead7..b201b3e7c57a6 100644 --- a/coderd/workspaceapps/db_test.go +++ b/coderd/workspaceapps/db_test.go @@ -171,7 +171,8 @@ func Test_ResolveRequest(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, firstUser.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, agentAuthToken).Wait(client, workspace.ID, agentName) + _ = agenttest.New(t, client.URL, agentAuthToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID, agentName) agentID := uuid.Nil for _, resource := range resources { diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index 90166e71d5e6f..f629ad773aa32 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -1347,7 +1347,8 @@ func TestWorkspaceFilterManual(t *testing.T) { workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -2258,13 +2259,13 @@ func TestWorkspaceWatcher(t *testing.T) { wait("agent timeout after start", nil) agt := agenttest.New(t, client.URL, authToken) - agt.Wait(client, workspace.ID) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) wait("agent connected/ready", func(w codersdk.Workspace) bool { return w.LatestBuild.Resources[0].Agents[0].Status == codersdk.WorkspaceAgentConnected && w.LatestBuild.Resources[0].Agents[0].LifecycleState == codersdk.WorkspaceAgentLifecycleReady }) - agt.Agent().Close() + agt.Close() wait("agent disconnected", func(w codersdk.Workspace) bool { return w.LatestBuild.Resources[0].Agents[0].Status == codersdk.WorkspaceAgentDisconnected }) diff --git a/enterprise/coderd/workspaceproxy_test.go b/enterprise/coderd/workspaceproxy_test.go index 40cd2592ba743..1a7b5ce71409c 100644 --- a/enterprise/coderd/workspaceproxy_test.go +++ b/enterprise/coderd/workspaceproxy_test.go @@ -674,10 +674,8 @@ func TestIssueSignedAppToken(t *testing.T) { workspace.LatestBuild = build // Connect an agent to the workspace - - _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) - - coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) createProxyCtx := testutil.Context(t, testutil.WaitLong) proxyRes, err := client.CreateWorkspaceProxy(createProxyCtx, codersdk.CreateWorkspaceProxyRequest{ @@ -777,7 +775,8 @@ func TestReconnectingPTYSignedToken(t *testing.T) { // Connect an agent to the workspace agentID := build.Resources[0].Agents[0].ID - _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) proxyURL, err := url.Parse(fmt.Sprintf("https://%s.com", namesgenerator.GetRandomName(1))) require.NoError(t, err) diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go index 18b1017df2036..be5c1eee7949d 100644 --- a/enterprise/wsproxy/wsproxy_test.go +++ b/enterprise/wsproxy/wsproxy_test.go @@ -183,7 +183,8 @@ resourceLoop: require.NotEqual(t, uuid.Nil, agentID) // Connect an agent to the workspace - _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) t.Run("ReturnedInDERPMap", func(t *testing.T) { t.Parallel() @@ -407,7 +408,8 @@ resourceLoop: require.NotEqual(t, uuid.Nil, agentID) // Connect an agent to the workspace - _ = agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + _ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) // Connect to the workspace agent. ctx := testutil.Context(t, testutil.WaitLong) diff --git a/scaletest/agentconn/run_test.go b/scaletest/agentconn/run_test.go index 36dae72ef6e82..70f373dc735de 100644 --- a/scaletest/agentconn/run_test.go +++ b/scaletest/agentconn/run_test.go @@ -256,7 +256,8 @@ func setupRunnerTest(t *testing.T) (client *codersdk.Client, agentID uuid.UUID) workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) return client, resources[0].Agents[0].ID } diff --git a/scaletest/reconnectingpty/run_test.go b/scaletest/reconnectingpty/run_test.go index b85df2437a1a9..407f1562b59bd 100644 --- a/scaletest/reconnectingpty/run_test.go +++ b/scaletest/reconnectingpty/run_test.go @@ -276,7 +276,8 @@ func setupRunnerTest(t *testing.T) (client *codersdk.Client, agentID uuid.UUID) workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - resources := agenttest.New(t, client.URL, authToken).Wait(client, workspace.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) require.Eventually(t, func() bool { t.Log("agent id", resources[0].Agents[0].ID) return (*api.TailnetCoordinator.Load()).Node(resources[0].Agents[0].ID) != nil diff --git a/scaletest/workspacetraffic/run_test.go b/scaletest/workspacetraffic/run_test.go index 604b606806569..043ed31cb8533 100644 --- a/scaletest/workspacetraffic/run_test.go +++ b/scaletest/workspacetraffic/run_test.go @@ -73,7 +73,8 @@ func TestRun(t *testing.T) { ) // We also need a running agent to run this test. - resources := agenttest.New(t, client.URL, authToken).Wait(client, ws.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID) ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) @@ -191,7 +192,8 @@ func TestRun(t *testing.T) { ) // We also need a running agent to run this test. - resources := agenttest.New(t, client.URL, authToken).Wait(client, ws.ID) + _ = agenttest.New(t, client.URL, authToken) + resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID) ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) From 903fe3eb47f6843f96edd807078f004649fe0829 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Mon, 25 Sep 2023 12:49:15 +0100 Subject: [PATCH 9/9] skip flaky test --- scaletest/dashboard/run_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/scaletest/dashboard/run_test.go b/scaletest/dashboard/run_test.go index 88c4aecb9b8b6..bd18359e19eed 100644 --- a/scaletest/dashboard/run_test.go +++ b/scaletest/dashboard/run_test.go @@ -19,6 +19,7 @@ import ( func Test_Run(t *testing.T) { t.Parallel() + t.Skip("To be fixed by https://github.com/coder/coder/issues/9131") if testutil.RaceEnabled() { t.Skip("skipping timing-sensitive test because of race detector") }