Skip to content

Commit 46dfec7

Browse files
spikecurtisspikecurtis
committed
chore: refactor CLI agent auth tests as unit tests
1 parent a13a334 commit 46dfec7

File tree

12 files changed

+121
-196
lines changed

12 files changed

+121
-196
lines changed

‎cli/agent.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
178178
slog.F("auth", r.agentAuth),
179179
slog.F("version", version),
180180
)
181-
client, err := r.createAgentClient(ctx)
181+
client, err := r.CreateAgentClient(ctx)
182182
if err != nil {
183183
return xerrors.Errorf("create agent client: %w", err)
184184
}

‎cli/agent_test.go

Lines changed: 0 additions & 157 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package cli_test
22

33
import (
4-
"context"
54
"fmt"
65
"net/http"
76
"os"
@@ -11,7 +10,6 @@ import (
1110
"sync/atomic"
1211
"testing"
1312

14-
"github.com/google/uuid"
1513
"github.com/stretchr/testify/assert"
1614
"github.com/stretchr/testify/require"
1715

@@ -21,10 +19,7 @@ import (
2119
"github.com/coder/coder/v2/coderd/coderdtest"
2220
"github.com/coder/coder/v2/coderd/database"
2321
"github.com/coder/coder/v2/coderd/database/dbfake"
24-
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2522
"github.com/coder/coder/v2/codersdk"
26-
"github.com/coder/coder/v2/codersdk/workspacesdk"
27-
"github.com/coder/coder/v2/provisionersdk/proto"
2823
"github.com/coder/coder/v2/testutil"
2924
)
3025

@@ -64,158 +59,6 @@ func TestWorkspaceAgent(t *testing.T) {
6459
}, testutil.WaitLong, testutil.IntervalMedium)
6560
})
6661

67-
t.Run("Azure", func(t *testing.T) {
68-
t.Parallel()
69-
instanceID := "instanceidentifier"
70-
certificates, metadataClient := coderdtest.NewAzureInstanceIdentity(t, instanceID)
71-
db, ps := dbtestutil.NewDB(t,
72-
dbtestutil.WithDumpOnFailure(),
73-
)
74-
client := coderdtest.New(t, &coderdtest.Options{
75-
Database: db,
76-
Pubsub: ps,
77-
AzureCertificates: certificates,
78-
})
79-
user := coderdtest.CreateFirstUser(t, client)
80-
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
81-
OrganizationID: user.OrganizationID,
82-
OwnerID: user.UserID,
83-
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
84-
agents[0].Auth = &proto.Agent_InstanceId{InstanceId: instanceID}
85-
return agents
86-
}).Do()
87-
88-
inv, _ := clitest.New(t, "agent", "--auth", "azure-instance-identity", "--agent-url", client.URL.String())
89-
inv = inv.WithContext(
90-
//nolint:revive,staticcheck
91-
context.WithValue(inv.Context(), "azure-client", metadataClient),
92-
)
93-
94-
ctx := inv.Context()
95-
clitest.Start(t, inv)
96-
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
97-
MatchResources(matchAgentWithVersion).Wait()
98-
workspace, err := client.Workspace(ctx, r.Workspace.ID)
99-
require.NoError(t, err)
100-
resources := workspace.LatestBuild.Resources
101-
if assert.NotEmpty(t, workspace.LatestBuild.Resources) && assert.NotEmpty(t, resources[0].Agents) {
102-
assert.NotEmpty(t, resources[0].Agents[0].Version)
103-
}
104-
dialer, err := workspacesdk.New(client).
105-
DialAgent(ctx, resources[0].Agents[0].ID, nil)
106-
require.NoError(t, err)
107-
defer dialer.Close()
108-
require.True(t, dialer.AwaitReachable(ctx))
109-
})
110-
111-
t.Run("AWS", func(t *testing.T) {
112-
t.Parallel()
113-
instanceID := "instanceidentifier"
114-
certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
115-
db, ps := dbtestutil.NewDB(t,
116-
dbtestutil.WithDumpOnFailure(),
117-
)
118-
client := coderdtest.New(t, &coderdtest.Options{
119-
Database: db,
120-
Pubsub: ps,
121-
AWSCertificates: certificates,
122-
})
123-
user := coderdtest.CreateFirstUser(t, client)
124-
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
125-
OrganizationID: user.OrganizationID,
126-
OwnerID: user.UserID,
127-
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
128-
agents[0].Auth = &proto.Agent_InstanceId{InstanceId: instanceID}
129-
return agents
130-
}).Do()
131-
132-
inv, _ := clitest.New(t, "agent", "--auth", "aws-instance-identity", "--agent-url", client.URL.String())
133-
inv = inv.WithContext(
134-
//nolint:revive,staticcheck
135-
context.WithValue(inv.Context(), "aws-client", metadataClient),
136-
)
137-
138-
clitest.Start(t, inv)
139-
ctx := inv.Context()
140-
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
141-
MatchResources(matchAgentWithVersion).
142-
Wait()
143-
workspace, err := client.Workspace(ctx, r.Workspace.ID)
144-
require.NoError(t, err)
145-
resources := workspace.LatestBuild.Resources
146-
if assert.NotEmpty(t, resources) && assert.NotEmpty(t, resources[0].Agents) {
147-
assert.NotEmpty(t, resources[0].Agents[0].Version)
148-
}
149-
dialer, err := workspacesdk.New(client).
150-
DialAgent(ctx, resources[0].Agents[0].ID, nil)
151-
require.NoError(t, err)
152-
defer dialer.Close()
153-
require.True(t, dialer.AwaitReachable(ctx))
154-
})
155-
156-
t.Run("GoogleCloud", func(t *testing.T) {
157-
t.Parallel()
158-
instanceID := "instanceidentifier"
159-
validator, metadataClient := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
160-
db, ps := dbtestutil.NewDB(t,
161-
dbtestutil.WithDumpOnFailure(),
162-
)
163-
client := coderdtest.New(t, &coderdtest.Options{
164-
Database: db,
165-
Pubsub: ps,
166-
GoogleTokenValidator: validator,
167-
})
168-
owner := coderdtest.CreateFirstUser(t, client)
169-
member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
170-
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
171-
OrganizationID: owner.OrganizationID,
172-
OwnerID: memberUser.ID,
173-
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
174-
agents[0].Auth = &proto.Agent_InstanceId{InstanceId: instanceID}
175-
return agents
176-
}).Do()
177-
178-
inv, cfg := clitest.New(t, "agent", "--auth", "google-instance-identity", "--agent-url", client.URL.String())
179-
clitest.SetupConfig(t, member, cfg)
180-
181-
clitest.Start(t,
182-
inv.WithContext(
183-
//nolint:revive,staticcheck
184-
context.WithValue(inv.Context(), "gcp-client", metadataClient),
185-
),
186-
)
187-
188-
ctx := inv.Context()
189-
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
190-
MatchResources(matchAgentWithVersion).
191-
Wait()
192-
workspace, err := client.Workspace(ctx, r.Workspace.ID)
193-
require.NoError(t, err)
194-
resources := workspace.LatestBuild.Resources
195-
if assert.NotEmpty(t, resources) && assert.NotEmpty(t, resources[0].Agents) {
196-
assert.NotEmpty(t, resources[0].Agents[0].Version)
197-
}
198-
dialer, err := workspacesdk.New(client).DialAgent(ctx, resources[0].Agents[0].ID, nil)
199-
require.NoError(t, err)
200-
defer dialer.Close()
201-
require.True(t, dialer.AwaitReachable(ctx))
202-
sshClient, err := dialer.SSHClient(ctx)
203-
require.NoError(t, err)
204-
defer sshClient.Close()
205-
session, err := sshClient.NewSession()
206-
require.NoError(t, err)
207-
defer session.Close()
208-
key := "CODER_AGENT_TOKEN"
209-
command := "sh -c 'echo $" + key + "'"
210-
if runtime.GOOS == "windows" {
211-
command = "cmd.exe /c echo %" + key + "%"
212-
}
213-
token, err := session.CombinedOutput(command)
214-
require.NoError(t, err)
215-
_, err = uuid.Parse(strings.TrimSpace(string(token)))
216-
require.NoError(t, err)
217-
})
218-
21962
t.Run("PostStartup", func(t *testing.T) {
22063
t.Parallel()
22164

‎cli/exp_mcp.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
148148
binPath = testBinaryName
149149
}
150150
configureClaudeEnv := map[string]string{}
151-
agentClient, err := r.createAgentClient(inv.Context())
151+
agentClient, err := r.CreateAgentClient(inv.Context())
152152
if err != nil {
153153
cliui.Warnf(inv.Stderr, "failed to create agent client: %s", err)
154154
} else {
@@ -494,7 +494,7 @@ func (r *RootCmd) mcpServer() *serpent.Command {
494494
}
495495

496496
// Try to create an agent client for status reporting. Not validated.
497-
agentClient, err := r.createAgentClient(inv.Context())
497+
agentClient, err := r.CreateAgentClient(inv.Context())
498498
if err == nil {
499499
cliui.Infof(inv.Stderr, "Agent URL : %s", agentClient.SDK.URL.String())
500500
srv.agentClient = agentClient

‎cli/externalauth.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ fi
7575
return xerrors.Errorf("agent token not found")
7676
}
7777

78-
client, err := r.createAgentClient(ctx)
78+
client, err := r.CreateAgentClient(ctx)
7979
if err != nil {
8080
return xerrors.Errorf("create agent client: %w", err)
8181
}

‎cli/gitaskpass.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func (r *RootCmd) gitAskpass() *serpent.Command {
3333
return xerrors.Errorf("parse host: %w", err)
3434
}
3535

36-
client, err := r.createAgentClient(ctx)
36+
client, err := r.CreateAgentClient(ctx)
3737
if err != nil {
3838
return xerrors.Errorf("create agent client: %w", err)
3939
}

‎cli/gitssh.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func (r *RootCmd) gitssh() *serpent.Command {
3838
return err
3939
}
4040

41-
client, err := r.createAgentClient(ctx)
41+
client, err := r.CreateAgentClient(ctx)
4242
if err != nil {
4343
return xerrors.Errorf("create agent client: %w", err)
4444
}

‎cli/root.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -688,9 +688,9 @@ func (r *RootCmd) createUnauthenticatedClient(ctx context.Context, serverURL *ur
688688
return &client, err
689689
}
690690

691-
// createAgentClient returns a new client from the command context. It works
691+
// CreateAgentClient returns a new client from the command context. It works
692692
// just like InitClient, but uses the agent token and URL instead.
693-
func (r *RootCmd) createAgentClient(ctx context.Context) (*agentsdk.Client, error) {
693+
func (r *RootCmd) CreateAgentClient(ctx context.Context) (*agentsdk.Client, error) {
694694
agentURL := r.agentURL
695695
if agentURL == nil || agentURL.String() == "" {
696696
return nil, xerrors.Errorf("%s must be set", envAgentURL)

‎cli/root_test.go

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,14 @@ import (
1010
"sync/atomic"
1111
"testing"
1212

13+
"golang.org/x/xerrors"
14+
1315
"github.com/coder/serpent"
1416

1517
"github.com/coder/coder/v2/coderd"
1618
"github.com/coder/coder/v2/coderd/coderdtest"
1719
"github.com/coder/coder/v2/codersdk"
20+
"github.com/coder/coder/v2/codersdk/agentsdk"
1821
"github.com/coder/coder/v2/pty/ptytest"
1922
"github.com/coder/coder/v2/testutil"
2023

@@ -275,3 +278,82 @@ func TestHandlersOK(t *testing.T) {
275278

276279
clitest.HandlersOK(t, cmd)
277280
}
281+
282+
func TestCreateAgentClient_Token(t *testing.T) {
283+
t.Parallel()
284+
285+
client := createAgentWithFlags(t,
286+
"--agent-token", "fake-token",
287+
"--agent-url", "http://coder.fake")
288+
require.Equal(t, "fake-token", client.GetSessionToken())
289+
}
290+
291+
func TestCreateAgentClient_Google(t *testing.T) {
292+
t.Parallel()
293+
294+
client := createAgentWithFlags(t,
295+
"--auth", "google-instance-identity",
296+
"--agent-url", "http://coder.fake")
297+
provider, ok := client.RefreshableSessionTokenProvider.(*agentsdk.InstanceIdentitySessionTokenProvider)
298+
require.True(t, ok)
299+
require.NotNil(t, provider.TokenExchanger)
300+
require.IsType(t, &agentsdk.GoogleSessionTokenExchanger{}, provider.TokenExchanger)
301+
}
302+
303+
func TestCreateAgentClient_AWS(t *testing.T) {
304+
t.Parallel()
305+
306+
client := createAgentWithFlags(t,
307+
"--auth", "aws-instance-identity",
308+
"--agent-url", "http://coder.fake")
309+
provider, ok := client.RefreshableSessionTokenProvider.(*agentsdk.InstanceIdentitySessionTokenProvider)
310+
require.True(t, ok)
311+
require.NotNil(t, provider.TokenExchanger)
312+
require.IsType(t, &agentsdk.AWSSessionTokenExchanger{}, provider.TokenExchanger)
313+
}
314+
315+
func TestCreateAgentClient_Azure(t *testing.T) {
316+
t.Parallel()
317+
318+
client := createAgentWithFlags(t,
319+
"--auth", "azure-instance-identity",
320+
"--agent-url", "http://coder.fake")
321+
provider, ok := client.RefreshableSessionTokenProvider.(*agentsdk.InstanceIdentitySessionTokenProvider)
322+
require.True(t, ok)
323+
require.NotNil(t, provider.TokenExchanger)
324+
require.IsType(t, &agentsdk.AzureSessionTokenExchanger{}, provider.TokenExchanger)
325+
}
326+
327+
func createAgentWithFlags(t *testing.T, flags ...string) *agentsdk.Client {
328+
t.Helper()
329+
r := &cli.RootCmd{}
330+
var client *agentsdk.Client
331+
subCmd := agentClientCommand(r, &client)
332+
cmd, err := r.Command([]*serpent.Command{subCmd})
333+
require.NoError(t, err)
334+
inv, _ := clitest.NewWithCommand(t, cmd,
335+
append([]string{"agent-client"}, flags...)...)
336+
err = inv.Run()
337+
require.NoError(t, err)
338+
require.NotNil(t, client)
339+
return client
340+
}
341+
342+
// agentClientCommand creates a subcommand that creates an agent client and stores it in the provided clientRef. Used to
343+
// test the properties of the client with various root command flags.
344+
func agentClientCommand(r *cli.RootCmd, clientRef **agentsdk.Client) *serpent.Command {
345+
return &serpent.Command{
346+
Use: "agent-client",
347+
Short: `Creates and agent client for testing.`,
348+
// This command isn't useful to manually execute.
349+
Hidden: true,
350+
Handler: func(inv *serpent.Invocation) error {
351+
client, err := r.CreateAgentClient(inv.Context())
352+
if err != nil {
353+
return xerrors.Errorf("create agent client: %w", err)
354+
}
355+
*clientRef = client
356+
return nil
357+
},
358+
}
359+
}

0 commit comments

Comments
 (0)