@@ -19,6 +19,7 @@ import (
19
19
"github.com/coder/coder/v2/coderd/database"
20
20
"github.com/coder/coder/v2/coderd/database/dbmock"
21
21
"github.com/coder/coder/v2/coderd/database/dbtime"
22
+ "github.com/coder/coder/v2/coderd/database/pubsub"
22
23
"github.com/coder/coder/v2/coderd/prometheusmetrics"
23
24
"github.com/coder/coder/v2/coderd/schedule"
24
25
)
@@ -78,8 +79,10 @@ func TestUpdateStates(t *testing.T) {
78
79
t .Parallel ()
79
80
80
81
var (
81
- now = dbtime .Now ()
82
- dbM = dbmock .NewMockStore (gomock .NewController (t ))
82
+ now = dbtime .Now ()
83
+ db = dbmock .NewMockStore (gomock .NewController (t ))
84
+ ps = pubsub .NewInMemory ()
85
+
83
86
templateScheduleStore = schedule.MockTemplateScheduleStore {
84
87
GetFn : func (context.Context , database.Store , uuid.UUID ) (schedule.TemplateScheduleOptions , error ) {
85
88
panic ("should not be called" )
@@ -124,7 +127,8 @@ func TestUpdateStates(t *testing.T) {
124
127
AgentFn : func (context.Context ) (database.WorkspaceAgent , error ) {
125
128
return agent , nil
126
129
},
127
- Database : dbM ,
130
+ Database : db ,
131
+ Pubsub : ps ,
128
132
StatsBatcher : batcher ,
129
133
TemplateScheduleStore : templateScheduleStorePtr (templateScheduleStore ),
130
134
AgentStatsRefreshInterval : 10 * time .Second ,
@@ -144,25 +148,25 @@ func TestUpdateStates(t *testing.T) {
144
148
}
145
149
146
150
// Workspace gets fetched.
147
- dbM .EXPECT ().GetWorkspaceByAgentID (gomock .Any (), agent .ID ).Return (database.GetWorkspaceByAgentIDRow {
151
+ db .EXPECT ().GetWorkspaceByAgentID (gomock .Any (), agent .ID ).Return (database.GetWorkspaceByAgentIDRow {
148
152
Workspace : workspace ,
149
153
TemplateName : template .Name ,
150
154
}, nil )
151
155
152
156
// We expect an activity bump because ConnectionCount > 0.
153
- dbM .EXPECT ().ActivityBumpWorkspace (gomock .Any (), database.ActivityBumpWorkspaceParams {
157
+ db .EXPECT ().ActivityBumpWorkspace (gomock .Any (), database.ActivityBumpWorkspaceParams {
154
158
WorkspaceID : workspace .ID ,
155
159
NextAutostart : time.Time {}.UTC (),
156
160
}).Return (nil )
157
161
158
162
// Workspace last used at gets bumped.
159
- dbM .EXPECT ().UpdateWorkspaceLastUsedAt (gomock .Any (), database.UpdateWorkspaceLastUsedAtParams {
163
+ db .EXPECT ().UpdateWorkspaceLastUsedAt (gomock .Any (), database.UpdateWorkspaceLastUsedAtParams {
160
164
ID : workspace .ID ,
161
165
LastUsedAt : now ,
162
166
}).Return (nil )
163
167
164
168
// User gets fetched to hit the UpdateAgentMetricsFn.
165
- dbM .EXPECT ().GetUserByID (gomock .Any (), user .ID ).Return (user , nil )
169
+ db .EXPECT ().GetUserByID (gomock .Any (), user .ID ).Return (user , nil )
166
170
167
171
resp , err := api .UpdateStats (context .Background (), req )
168
172
require .NoError (t , err )
@@ -188,7 +192,8 @@ func TestUpdateStates(t *testing.T) {
188
192
189
193
var (
190
194
now = dbtime .Now ()
191
- dbM = dbmock .NewMockStore (gomock .NewController (t ))
195
+ db = dbmock .NewMockStore (gomock .NewController (t ))
196
+ ps = pubsub .NewInMemory ()
192
197
templateScheduleStore = schedule.MockTemplateScheduleStore {
193
198
GetFn : func (context.Context , database.Store , uuid.UUID ) (schedule.TemplateScheduleOptions , error ) {
194
199
panic ("should not be called" )
@@ -213,7 +218,8 @@ func TestUpdateStates(t *testing.T) {
213
218
AgentFn : func (context.Context ) (database.WorkspaceAgent , error ) {
214
219
return agent , nil
215
220
},
216
- Database : dbM ,
221
+ Database : db ,
222
+ Pubsub : ps ,
217
223
StatsBatcher : batcher ,
218
224
TemplateScheduleStore : templateScheduleStorePtr (templateScheduleStore ),
219
225
AgentStatsRefreshInterval : 10 * time .Second ,
@@ -225,13 +231,13 @@ func TestUpdateStates(t *testing.T) {
225
231
}
226
232
227
233
// Workspace gets fetched.
228
- dbM .EXPECT ().GetWorkspaceByAgentID (gomock .Any (), agent .ID ).Return (database.GetWorkspaceByAgentIDRow {
234
+ db .EXPECT ().GetWorkspaceByAgentID (gomock .Any (), agent .ID ).Return (database.GetWorkspaceByAgentIDRow {
229
235
Workspace : workspace ,
230
236
TemplateName : template .Name ,
231
237
}, nil )
232
238
233
239
// Workspace last used at gets bumped.
234
- dbM .EXPECT ().UpdateWorkspaceLastUsedAt (gomock .Any (), database.UpdateWorkspaceLastUsedAtParams {
240
+ db .EXPECT ().UpdateWorkspaceLastUsedAt (gomock .Any (), database.UpdateWorkspaceLastUsedAtParams {
235
241
ID : workspace .ID ,
236
242
LastUsedAt : now ,
237
243
}).Return (nil )
@@ -244,7 +250,8 @@ func TestUpdateStates(t *testing.T) {
244
250
t .Parallel ()
245
251
246
252
var (
247
- dbM = dbmock .NewMockStore (gomock .NewController (t ))
253
+ db = dbmock .NewMockStore (gomock .NewController (t ))
254
+ ps = pubsub .NewInMemory ()
248
255
req = & agentproto.UpdateStatsRequest {
249
256
Stats : & agentproto.Stats {
250
257
ConnectionsByProto : map [string ]int64 {}, // len() == 0
@@ -255,7 +262,8 @@ func TestUpdateStates(t *testing.T) {
255
262
AgentFn : func (context.Context ) (database.WorkspaceAgent , error ) {
256
263
return agent , nil
257
264
},
258
- Database : dbM ,
265
+ Database : db ,
266
+ Pubsub : ps ,
259
267
StatsBatcher : nil , // should not be called
260
268
TemplateScheduleStore : nil , // should not be called
261
269
AgentStatsRefreshInterval : 10 * time .Second ,
@@ -290,7 +298,8 @@ func TestUpdateStates(t *testing.T) {
290
298
nextAutostart := now .Add (30 * time .Minute ).UTC () // always sent to DB as UTC
291
299
292
300
var (
293
- dbM = dbmock .NewMockStore (gomock .NewController (t ))
301
+ db = dbmock .NewMockStore (gomock .NewController (t ))
302
+ ps = pubsub .NewInMemory ()
294
303
templateScheduleStore = schedule.MockTemplateScheduleStore {
295
304
GetFn : func (context.Context , database.Store , uuid.UUID ) (schedule.TemplateScheduleOptions , error ) {
296
305
return schedule.TemplateScheduleOptions {
@@ -321,7 +330,8 @@ func TestUpdateStates(t *testing.T) {
321
330
AgentFn : func (context.Context ) (database.WorkspaceAgent , error ) {
322
331
return agent , nil
323
332
},
324
- Database : dbM ,
333
+ Database : db ,
334
+ Pubsub : ps ,
325
335
StatsBatcher : batcher ,
326
336
TemplateScheduleStore : templateScheduleStorePtr (templateScheduleStore ),
327
337
AgentStatsRefreshInterval : 15 * time .Second ,
@@ -341,26 +351,26 @@ func TestUpdateStates(t *testing.T) {
341
351
}
342
352
343
353
// Workspace gets fetched.
344
- dbM .EXPECT ().GetWorkspaceByAgentID (gomock .Any (), agent .ID ).Return (database.GetWorkspaceByAgentIDRow {
354
+ db .EXPECT ().GetWorkspaceByAgentID (gomock .Any (), agent .ID ).Return (database.GetWorkspaceByAgentIDRow {
345
355
Workspace : workspace ,
346
356
TemplateName : template .Name ,
347
357
}, nil )
348
358
349
359
// We expect an activity bump because ConnectionCount > 0. However, the
350
360
// next autostart time will be set on the bump.
351
- dbM .EXPECT ().ActivityBumpWorkspace (gomock .Any (), database.ActivityBumpWorkspaceParams {
361
+ db .EXPECT ().ActivityBumpWorkspace (gomock .Any (), database.ActivityBumpWorkspaceParams {
352
362
WorkspaceID : workspace .ID ,
353
363
NextAutostart : nextAutostart ,
354
364
}).Return (nil )
355
365
356
366
// Workspace last used at gets bumped.
357
- dbM .EXPECT ().UpdateWorkspaceLastUsedAt (gomock .Any (), database.UpdateWorkspaceLastUsedAtParams {
367
+ db .EXPECT ().UpdateWorkspaceLastUsedAt (gomock .Any (), database.UpdateWorkspaceLastUsedAtParams {
358
368
ID : workspace .ID ,
359
369
LastUsedAt : now ,
360
370
}).Return (nil )
361
371
362
372
// User gets fetched to hit the UpdateAgentMetricsFn.
363
- dbM .EXPECT ().GetUserByID (gomock .Any (), user .ID ).Return (user , nil )
373
+ db .EXPECT ().GetUserByID (gomock .Any (), user .ID ).Return (user , nil )
364
374
365
375
resp , err := api .UpdateStats (context .Background (), req )
366
376
require .NoError (t , err )
0 commit comments