Skip to content

Commit 3248673

Browse files
f0sself0ssel
committed
add ssh and tests
1 parent 1ab57a4 commit 3248673

File tree

6 files changed

+150
-47
lines changed

6 files changed

+150
-47
lines changed

‎cli/ssh.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,13 @@ func (r *RootCmd) ssh() *serpent.Command {
408408
return xerrors.Errorf("start shell: %w", err)
409409
}
410410

411+
// track workspace usage while connection is open
412+
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
413+
AgentID: workspaceAgent.ID,
414+
AppName: codersdk.UsageAppNameSSH,
415+
})
416+
defer closeUsage()
417+
411418
// Put cancel at the top of the defer stack to initiate
412419
// shutdown of services.
413420
defer cancel()

‎cli/ssh_test.go

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,13 +36,15 @@ import (
3636
"github.com/coder/coder/v2/agent"
3737
"github.com/coder/coder/v2/agent/agentssh"
3838
"github.com/coder/coder/v2/agent/agenttest"
39+
agentproto "github.com/coder/coder/v2/agent/proto"
3940
"github.com/coder/coder/v2/cli/clitest"
4041
"github.com/coder/coder/v2/cli/cliui"
4142
"github.com/coder/coder/v2/coderd/coderdtest"
4243
"github.com/coder/coder/v2/coderd/database"
4344
"github.com/coder/coder/v2/coderd/database/dbfake"
4445
"github.com/coder/coder/v2/coderd/database/dbtestutil"
4546
"github.com/coder/coder/v2/coderd/rbac"
47+
"github.com/coder/coder/v2/coderd/workspacestats/wstest"
4648
"github.com/coder/coder/v2/codersdk"
4749
"github.com/coder/coder/v2/provisioner/echo"
4850
"github.com/coder/coder/v2/provisionersdk/proto"
@@ -1292,6 +1294,89 @@ func TestSSH(t *testing.T) {
12921294
require.NoError(t, err)
12931295
require.Len(t, ents, 1, "expected one file in logdir %s", logDir)
12941296
})
1297+
t.Run("UpdateUsageNoExperiment", func(t *testing.T) {
1298+
t.Parallel()
1299+
1300+
batcher := &wstest.StatsBatcher{
1301+
LastStats: &agentproto.Stats{},
1302+
}
1303+
admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
1304+
StatsBatcher: batcher,
1305+
})
1306+
admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
1307+
first := coderdtest.CreateFirstUser(t, admin)
1308+
client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID)
1309+
r := dbfake.WorkspaceBuild(t, store, database.Workspace{
1310+
OrganizationID: first.OrganizationID,
1311+
OwnerID: user.ID,
1312+
}).WithAgent().Do()
1313+
workspace := r.Workspace
1314+
agentToken := r.AgentToken
1315+
inv, root := clitest.New(t, "ssh", workspace.Name)
1316+
clitest.SetupConfig(t, client, root)
1317+
pty := ptytest.New(t).Attach(inv)
1318+
1319+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
1320+
defer cancel()
1321+
1322+
cmdDone := tGo(t, func() {
1323+
err := inv.WithContext(ctx).Run()
1324+
assert.NoError(t, err)
1325+
})
1326+
pty.ExpectMatch("Waiting")
1327+
1328+
_ = agenttest.New(t, client.URL, agentToken)
1329+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
1330+
1331+
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
1332+
pty.WriteLine("exit")
1333+
<-cmdDone
1334+
1335+
require.EqualValues(t, 0, batcher.Called)
1336+
require.EqualValues(t, 0, batcher.LastStats.SessionCountSsh)
1337+
})
1338+
t.Run("UpdateUsageExperiment", func(t *testing.T) {
1339+
t.Parallel()
1340+
1341+
dv := coderdtest.DeploymentValues(t)
1342+
dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)}
1343+
batcher := &wstest.StatsBatcher{}
1344+
admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
1345+
DeploymentValues: dv,
1346+
StatsBatcher: batcher,
1347+
})
1348+
admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
1349+
first := coderdtest.CreateFirstUser(t, admin)
1350+
client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID)
1351+
r := dbfake.WorkspaceBuild(t, store, database.Workspace{
1352+
OrganizationID: first.OrganizationID,
1353+
OwnerID: user.ID,
1354+
}).WithAgent().Do()
1355+
workspace := r.Workspace
1356+
agentToken := r.AgentToken
1357+
inv, root := clitest.New(t, "ssh", workspace.Name)
1358+
clitest.SetupConfig(t, client, root)
1359+
pty := ptytest.New(t).Attach(inv)
1360+
1361+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
1362+
defer cancel()
1363+
1364+
cmdDone := tGo(t, func() {
1365+
err := inv.WithContext(ctx).Run()
1366+
assert.NoError(t, err)
1367+
})
1368+
pty.ExpectMatch("Waiting")
1369+
1370+
_ = agenttest.New(t, client.URL, agentToken)
1371+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
1372+
1373+
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
1374+
pty.WriteLine("exit")
1375+
<-cmdDone
1376+
1377+
require.EqualValues(t, 1, batcher.Called)
1378+
require.EqualValues(t, 1, batcher.LastStats.SessionCountSsh)
1379+
})
12951380
}
12961381

12971382
//nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel.

‎coderd/agentapi/stats_test.go

Lines changed: 18 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package agentapi_test
33
import (
44
"context"
55
"database/sql"
6-
"sync"
76
"sync/atomic"
87
"testing"
98
"time"
@@ -23,37 +22,11 @@ import (
2322
"github.com/coder/coder/v2/coderd/prometheusmetrics"
2423
"github.com/coder/coder/v2/coderd/schedule"
2524
"github.com/coder/coder/v2/coderd/workspacestats"
25+
"github.com/coder/coder/v2/coderd/workspacestats/wstest"
2626
"github.com/coder/coder/v2/codersdk"
2727
"github.com/coder/coder/v2/testutil"
2828
)
2929

30-
type statsBatcher struct {
31-
mu sync.Mutex
32-
33-
called int64
34-
lastTime time.Time
35-
lastAgentID uuid.UUID
36-
lastTemplateID uuid.UUID
37-
lastUserID uuid.UUID
38-
lastWorkspaceID uuid.UUID
39-
lastStats *agentproto.Stats
40-
}
41-
42-
var _ workspacestats.Batcher = &statsBatcher{}
43-
44-
func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error {
45-
b.mu.Lock()
46-
defer b.mu.Unlock()
47-
b.called++
48-
b.lastTime = now
49-
b.lastAgentID = agentID
50-
b.lastTemplateID = templateID
51-
b.lastUserID = userID
52-
b.lastWorkspaceID = workspaceID
53-
b.lastStats = st
54-
return nil
55-
}
56-
5730
func TestUpdateStates(t *testing.T) {
5831
t.Parallel()
5932

@@ -94,7 +67,7 @@ func TestUpdateStates(t *testing.T) {
9467
panic("not implemented")
9568
},
9669
}
97-
batcher = &statsBatcher{}
70+
batcher = &wstest.StatsBatcher{}
9871
updateAgentMetricsFnCalled = false
9972

10073
req = &agentproto.UpdateStatsRequest{
@@ -188,15 +161,15 @@ func TestUpdateStates(t *testing.T) {
188161
ReportInterval: durationpb.New(10 * time.Second),
189162
}, resp)
190163

191-
batcher.mu.Lock()
192-
defer batcher.mu.Unlock()
193-
require.Equal(t, int64(1), batcher.called)
194-
require.Equal(t, now, batcher.lastTime)
195-
require.Equal(t, agent.ID, batcher.lastAgentID)
196-
require.Equal(t, template.ID, batcher.lastTemplateID)
197-
require.Equal(t, user.ID, batcher.lastUserID)
198-
require.Equal(t, workspace.ID, batcher.lastWorkspaceID)
199-
require.Equal(t, req.Stats, batcher.lastStats)
164+
batcher.Mu.Lock()
165+
defer batcher.Mu.Unlock()
166+
require.Equal(t, int64(1), batcher.Called)
167+
require.Equal(t, now, batcher.LastTime)
168+
require.Equal(t, agent.ID, batcher.LastAgentID)
169+
require.Equal(t, template.ID, batcher.LastTemplateID)
170+
require.Equal(t, user.ID, batcher.LastUserID)
171+
require.Equal(t, workspace.ID, batcher.LastWorkspaceID)
172+
require.Equal(t, req.Stats, batcher.LastStats)
200173
ctx := testutil.Context(t, testutil.WaitShort)
201174
select {
202175
case <-ctx.Done():
@@ -222,7 +195,7 @@ func TestUpdateStates(t *testing.T) {
222195
panic("not implemented")
223196
},
224197
}
225-
batcher = &statsBatcher{}
198+
batcher = &wstest.StatsBatcher{}
226199

227200
req = &agentproto.UpdateStatsRequest{
228201
Stats: &agentproto.Stats{
@@ -336,7 +309,7 @@ func TestUpdateStates(t *testing.T) {
336309
panic("not implemented")
337310
},
338311
}
339-
batcher = &statsBatcher{}
312+
batcher = &wstest.StatsBatcher{}
340313
updateAgentMetricsFnCalled = false
341314

342315
req = &agentproto.UpdateStatsRequest{
@@ -423,7 +396,7 @@ func TestUpdateStates(t *testing.T) {
423396
panic("not implemented")
424397
},
425398
}
426-
batcher = &statsBatcher{}
399+
batcher = &wstest.StatsBatcher{}
427400

428401
req = &agentproto.UpdateStatsRequest{
429402
Stats: &agentproto.Stats{
@@ -496,10 +469,10 @@ func TestUpdateStates(t *testing.T) {
496469
ReportInterval: durationpb.New(10 * time.Second),
497470
}, resp)
498471

499-
batcher.mu.Lock()
500-
defer batcher.mu.Unlock()
501-
require.EqualValues(t, int64(1), batcher.called)
502-
require.EqualValues(t, batcher.lastStats.SessionCountSsh, 0)
472+
batcher.Mu.Lock()
473+
defer batcher.Mu.Unlock()
474+
require.EqualValues(t, 1, batcher.Called)
475+
require.EqualValues(t, 0, batcher.LastStats.SessionCountSsh)
503476
// TODO: other session values will come as they are migrated over
504477
// require.EqualValues(t, batcher.lastStats.SessionCountVscode, 0)
505478
// require.EqualValues(t, batcher.lastStats.SessionCountJetbrains, 0)

‎coderd/coderd.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ type Options struct {
187187
HTTPClient *http.Client
188188

189189
UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
190-
StatsBatcher *workspacestats.DBBatcher
190+
StatsBatcher workspacestats.Batcher
191191

192192
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
193193

‎coderd/coderdtest/coderdtest.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ type Options struct {
145145
// Logger should only be overridden if you expect errors
146146
// as part of your test.
147147
Logger *slog.Logger
148-
StatsBatcher *workspacestats.DBBatcher
148+
StatsBatcher workspacestats.Batcher
149149

150150
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
151151
AllowWorkspaceRenames bool

‎coderd/workspacestats/wstest/batcher.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package wstest
2+
3+
import (
4+
"sync"
5+
"time"
6+
7+
"github.com/google/uuid"
8+
9+
agentproto "github.com/coder/coder/v2/agent/proto"
10+
"github.com/coder/coder/v2/coderd/workspacestats"
11+
)
12+
13+
type StatsBatcher struct {
14+
Mu sync.Mutex
15+
16+
Called int64
17+
LastTime time.Time
18+
LastAgentID uuid.UUID
19+
LastTemplateID uuid.UUID
20+
LastUserID uuid.UUID
21+
LastWorkspaceID uuid.UUID
22+
LastStats *agentproto.Stats
23+
}
24+
25+
var _ workspacestats.Batcher = &StatsBatcher{}
26+
27+
func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error {
28+
b.Mu.Lock()
29+
defer b.Mu.Unlock()
30+
b.Called++
31+
b.LastTime = now
32+
b.LastAgentID = agentID
33+
b.LastTemplateID = templateID
34+
b.LastUserID = userID
35+
b.LastWorkspaceID = workspaceID
36+
b.LastStats = st
37+
return nil
38+
}

0 commit comments

Comments
 (0)