Skip to content

Commit 7412e77

Browse files
f0sself0ssel
committed
fix tests
1 parent 12b1ccf commit 7412e77

File tree

1 file changed

+44
-8
lines changed

1 file changed

+44
-8
lines changed

‎coderd/agentapi/stats_test.go

Lines changed: 44 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,11 @@ func TestUpdateStates(t *testing.T) {
6969
}
7070
batcher = &workspacestatstest.StatsBatcher{}
7171
updateAgentMetricsFnCalled = false
72+
tickCh = make(chan time.Time)
73+
flushCh = make(chan int, 1)
74+
wut = workspacestats.NewTracker(dbM,
75+
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
76+
)
7277

7378
req = &agentproto.UpdateStatsRequest{
7479
Stats: &agentproto.Stats{
@@ -108,7 +113,7 @@ func TestUpdateStates(t *testing.T) {
108113
Database: dbM,
109114
Pubsub: ps,
110115
StatsBatcher: batcher,
111-
UsageTracker: workspacestats.NewTracker(dbM),
116+
UsageTracker: wut,
112117
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
113118
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
114119
updateAgentMetricsFnCalled = true
@@ -126,6 +131,7 @@ func TestUpdateStates(t *testing.T) {
126131
return now
127132
},
128133
}
134+
defer wut.Close()
129135

130136
// Workspace gets fetched.
131137
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
@@ -142,6 +148,12 @@ func TestUpdateStates(t *testing.T) {
142148
NextAutostart: time.Time{}.UTC(),
143149
}).Return(nil)
144150

151+
// Workspace last used at gets bumped.
152+
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
153+
IDs: []uuid.UUID{workspace.ID},
154+
LastUsedAt: now,
155+
}).Return(nil)
156+
145157
// Ensure that pubsub notifications are sent.
146158
notifyDescription := make(chan []byte)
147159
ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
@@ -156,6 +168,10 @@ func TestUpdateStates(t *testing.T) {
156168
ReportInterval: durationpb.New(10 * time.Second),
157169
}, resp)
158170

171+
tickCh <- now
172+
count := <-flushCh
173+
require.Equal(t, 1, count, "expected one flush with one id")
174+
159175
batcher.Mu.Lock()
160176
defer batcher.Mu.Unlock()
161177
require.Equal(t, int64(1), batcher.Called)
@@ -301,6 +317,11 @@ func TestUpdateStates(t *testing.T) {
301317
}
302318
batcher = &workspacestatstest.StatsBatcher{}
303319
updateAgentMetricsFnCalled = false
320+
tickCh = make(chan time.Time)
321+
flushCh = make(chan int, 1)
322+
wut = workspacestats.NewTracker(dbM,
323+
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
324+
)
304325

305326
req = &agentproto.UpdateStatsRequest{
306327
Stats: &agentproto.Stats{
@@ -320,7 +341,7 @@ func TestUpdateStates(t *testing.T) {
320341
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
321342
Database: dbM,
322343
Pubsub: ps,
323-
UsageTracker: workspacestats.NewTracker(dbM),
344+
UsageTracker: wut,
324345
StatsBatcher: batcher,
325346
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
326347
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
@@ -339,6 +360,7 @@ func TestUpdateStates(t *testing.T) {
339360
return now
340361
},
341362
}
363+
defer wut.Close()
342364

343365
// Workspace gets fetched.
344366
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
@@ -354,9 +376,9 @@ func TestUpdateStates(t *testing.T) {
354376
}).Return(nil)
355377

356378
// Workspace last used at gets bumped.
357-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
358-
ID: workspace.ID,
359-
LastUsedAt: now,
379+
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
380+
IDs: []uuid.UUID{workspace.ID},
381+
LastUsedAt: now.UTC(),
360382
}).Return(nil)
361383

362384
// User gets fetched to hit the UpdateAgentMetricsFn.
@@ -368,6 +390,10 @@ func TestUpdateStates(t *testing.T) {
368390
ReportInterval: durationpb.New(15 * time.Second),
369391
}, resp)
370392

393+
tickCh <- now
394+
count := <-flushCh
395+
require.Equal(t, 1, count, "expected one flush with one id")
396+
371397
require.True(t, updateAgentMetricsFnCalled)
372398
})
373399

@@ -391,6 +417,11 @@ func TestUpdateStates(t *testing.T) {
391417
}
392418
batcher = &workspacestatstest.StatsBatcher{}
393419
updateAgentMetricsFnCalled = false
420+
tickCh = make(chan time.Time)
421+
flushCh = make(chan int, 1)
422+
wut = workspacestats.NewTracker(dbM,
423+
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
424+
)
394425

395426
req = &agentproto.UpdateStatsRequest{
396427
Stats: &agentproto.Stats{
@@ -421,6 +452,7 @@ func TestUpdateStates(t *testing.T) {
421452
},
422453
}
423454
)
455+
defer wut.Close()
424456
api := agentapi.StatsAPI{
425457
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
426458
return agent, nil
@@ -430,7 +462,7 @@ func TestUpdateStates(t *testing.T) {
430462
Database: dbM,
431463
Pubsub: ps,
432464
StatsBatcher: batcher,
433-
UsageTracker: workspacestats.NewTracker(dbM),
465+
UsageTracker: wut,
434466
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
435467
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
436468
updateAgentMetricsFnCalled = true
@@ -465,8 +497,8 @@ func TestUpdateStates(t *testing.T) {
465497
}).Return(nil)
466498

467499
// Workspace last used at gets bumped.
468-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
469-
ID: workspace.ID,
500+
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
501+
IDs: []uuid.UUID{workspace.ID},
470502
LastUsedAt: now,
471503
}).Return(nil)
472504

@@ -487,6 +519,10 @@ func TestUpdateStates(t *testing.T) {
487519
ReportInterval: durationpb.New(10 * time.Second),
488520
}, resp)
489521

522+
tickCh <- now
523+
count := <-flushCh
524+
require.Equal(t, 1, count, "expected one flush with one id")
525+
490526
batcher.Mu.Lock()
491527
defer batcher.Mu.Unlock()
492528
require.EqualValues(t, 1, batcher.Called)

0 commit comments

Comments
 (0)