Skip to content

Commit 3fbbf99

Browse files
spikecurtisspikecurtis
committed
chore: refactor CLI agent auth tests as unit tests
1 parent e5bf69b commit 3fbbf99

File tree

12 files changed

+127
-238
lines changed

12 files changed

+127
-238
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()
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()
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()
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()
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()
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()
4242
if err != nil {
4343
return xerrors.Errorf("create agent client: %w", err)
4444
}

‎cli/root.go

Lines changed: 5 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import (
2424
"text/tabwriter"
2525
"time"
2626

27-
"cloud.google.com/go/compute/metadata"
2827
"github.com/mattn/go-isatty"
2928
"github.com/mitchellh/go-wordwrap"
3029
"golang.org/x/mod/semver"
@@ -688,9 +687,9 @@ func (r *RootCmd) createUnauthenticatedClient(ctx context.Context, serverURL *ur
688687
return &client, err
689688
}
690689

691-
// createAgentClient returns a new client from the command context. It works
690+
// CreateAgentClient returns a new client from the command context. It works
692691
// just like InitClient, but uses the agent token and URL instead.
693-
func (r *RootCmd) createAgentClient(ctx context.Context) (*agentsdk.Client, error) {
692+
func (r *RootCmd) CreateAgentClient() (*agentsdk.Client, error) {
694693
agentURL := r.agentURL
695694
if agentURL == nil || agentURL.String() == "" {
696695
return nil, xerrors.Errorf("%s must be set", envAgentURL)
@@ -714,41 +713,11 @@ func (r *RootCmd) createAgentClient(ctx context.Context) (*agentsdk.Client, erro
714713
}
715714
return agentsdk.New(r.agentURL, agentsdk.UsingFixedToken(token)), nil
716715
case "google-instance-identity":
717-
718-
// This is *only* done for testing to mock client authentication.
719-
// This will never be set in a production scenario.
720-
var gcpClient *metadata.Client
721-
gcpClientRaw := ctx.Value("gcp-client")
722-
if gcpClientRaw != nil {
723-
gcpClient, _ = gcpClientRaw.(*metadata.Client)
724-
}
725-
return agentsdk.New(r.agentURL, agentsdk.UsingGoogleInstanceIdentity("", gcpClient)), nil
716+
return agentsdk.New(r.agentURL, agentsdk.UsingGoogleInstanceIdentity("", nil)), nil
726717
case "aws-instance-identity":
727-
client := agentsdk.New(r.agentURL, agentsdk.UsingAWSInstanceIdentity())
728-
// This is *only* done for testing to mock client authentication.
729-
// This will never be set in a production scenario.
730-
var awsClient *http.Client
731-
awsClientRaw := ctx.Value("aws-client")
732-
if awsClientRaw != nil {
733-
awsClient, _ = awsClientRaw.(*http.Client)
734-
if awsClient != nil {
735-
client.SDK.HTTPClient = awsClient
736-
}
737-
}
738-
return client, nil
718+
return agentsdk.New(r.agentURL, agentsdk.UsingAWSInstanceIdentity()), nil
739719
case "azure-instance-identity":
740-
client := agentsdk.New(r.agentURL, agentsdk.UsingAzureInstanceIdentity())
741-
// This is *only* done for testing to mock client authentication.
742-
// This will never be set in a production scenario.
743-
var azureClient *http.Client
744-
azureClientRaw := ctx.Value("azure-client")
745-
if azureClientRaw != nil {
746-
azureClient, _ = azureClientRaw.(*http.Client)
747-
if azureClient != nil {
748-
client.SDK.HTTPClient = azureClient
749-
}
750-
}
751-
return client, nil
720+
return agentsdk.New(r.agentURL, agentsdk.UsingAzureInstanceIdentity()), nil
752721
default:
753722
return nil, xerrors.Errorf("unknown agent auth type: %s", r.agentAuth)
754723
}

0 commit comments

Comments
 (0)