@@ -353,6 +353,126 @@ func TestGetEligibleProvisionerDaemonsByProvisionerJobIDs(t *testing.T) {
353
353
})
354
354
}
355
355
356
+ func TestGetProvisionerDaemonsWithStatusByOrganization (t * testing.T ) {
357
+ t .Parallel ()
358
+
359
+ t .Run ("NoDaemonsInOrgReturnsEmpty" ,func (t * testing.T ) {
360
+ t .Parallel ()
361
+ db ,_ := dbtestutil .NewDB (t )
362
+ org := dbgen .Organization (t ,db , database.Organization {})
363
+ otherOrg := dbgen .Organization (t ,db , database.Organization {})
364
+ dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
365
+ Name :"non-matching-daemon" ,
366
+ OrganizationID :otherOrg .ID ,
367
+ })
368
+ daemons ,err := db .GetProvisionerDaemonsWithStatusByOrganization (context .Background (), database.GetProvisionerDaemonsWithStatusByOrganizationParams {
369
+ OrganizationID :org .ID ,
370
+ })
371
+ require .NoError (t ,err )
372
+ require .Empty (t ,daemons )
373
+ })
374
+
375
+ t .Run ("MatchesProvisionerIDs" ,func (t * testing.T ) {
376
+ t .Parallel ()
377
+ db ,_ := dbtestutil .NewDB (t )
378
+ org := dbgen .Organization (t ,db , database.Organization {})
379
+
380
+ matchingDaemon0 := dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
381
+ Name :"matching-daemon0" ,
382
+ OrganizationID :org .ID ,
383
+ })
384
+ matchingDaemon1 := dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
385
+ Name :"matching-daemon1" ,
386
+ OrganizationID :org .ID ,
387
+ })
388
+ dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
389
+ Name :"non-matching-daemon" ,
390
+ OrganizationID :org .ID ,
391
+ })
392
+
393
+ daemons ,err := db .GetProvisionerDaemonsWithStatusByOrganization (context .Background (), database.GetProvisionerDaemonsWithStatusByOrganizationParams {
394
+ OrganizationID :org .ID ,
395
+ IDs : []uuid.UUID {matchingDaemon0 .ID ,matchingDaemon1 .ID },
396
+ })
397
+ require .NoError (t ,err )
398
+ require .Len (t ,daemons ,2 )
399
+ if daemons [0 ].ProvisionerDaemon .ID != matchingDaemon0 .ID {
400
+ daemons [0 ],daemons [1 ]= daemons [1 ],daemons [0 ]
401
+ }
402
+ require .Equal (t ,matchingDaemon0 .ID ,daemons [0 ].ProvisionerDaemon .ID )
403
+ require .Equal (t ,matchingDaemon1 .ID ,daemons [1 ].ProvisionerDaemon .ID )
404
+ })
405
+
406
+ t .Run ("MatchesTags" ,func (t * testing.T ) {
407
+ t .Parallel ()
408
+ db ,_ := dbtestutil .NewDB (t )
409
+ org := dbgen .Organization (t ,db , database.Organization {})
410
+
411
+ fooDaemon := dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
412
+ Name :"foo-daemon" ,
413
+ OrganizationID :org .ID ,
414
+ Tags : database.StringMap {
415
+ "foo" :"bar" ,
416
+ },
417
+ })
418
+ dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
419
+ Name :"baz-daemon" ,
420
+ OrganizationID :org .ID ,
421
+ Tags : database.StringMap {
422
+ "baz" :"qux" ,
423
+ },
424
+ })
425
+
426
+ daemons ,err := db .GetProvisionerDaemonsWithStatusByOrganization (context .Background (), database.GetProvisionerDaemonsWithStatusByOrganizationParams {
427
+ OrganizationID :org .ID ,
428
+ Tags : database.StringMap {"foo" :"bar" },
429
+ })
430
+ require .NoError (t ,err )
431
+ require .Len (t ,daemons ,1 )
432
+ require .Equal (t ,fooDaemon .ID ,daemons [0 ].ProvisionerDaemon .ID )
433
+ })
434
+
435
+ t .Run ("UsesStaleInterval" ,func (t * testing.T ) {
436
+ t .Parallel ()
437
+ db ,_ := dbtestutil .NewDB (t )
438
+ org := dbgen .Organization (t ,db , database.Organization {})
439
+
440
+ daemon1 := dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
441
+ Name :"stale-daemon" ,
442
+ OrganizationID :org .ID ,
443
+ CreatedAt :dbtime .Now ().Add (- time .Hour ),
444
+ LastSeenAt : sql.NullTime {
445
+ Valid :true ,
446
+ Time :dbtime .Now ().Add (- time .Hour ),
447
+ },
448
+ })
449
+ daemon2 := dbgen .ProvisionerDaemon (t ,db , database.ProvisionerDaemon {
450
+ Name :"idle-daemon" ,
451
+ OrganizationID :org .ID ,
452
+ CreatedAt :dbtime .Now ().Add (- (30 * time .Minute )),
453
+ LastSeenAt : sql.NullTime {
454
+ Valid :true ,
455
+ Time :dbtime .Now ().Add (- (30 * time .Minute )),
456
+ },
457
+ })
458
+
459
+ daemons ,err := db .GetProvisionerDaemonsWithStatusByOrganization (context .Background (), database.GetProvisionerDaemonsWithStatusByOrganizationParams {
460
+ OrganizationID :org .ID ,
461
+ StaleIntervalMS :45 * time .Minute .Milliseconds (),
462
+ })
463
+ require .NoError (t ,err )
464
+ require .Len (t ,daemons ,2 )
465
+
466
+ if daemons [0 ].ProvisionerDaemon .ID != daemon1 .ID {
467
+ daemons [0 ],daemons [1 ]= daemons [1 ],daemons [0 ]
468
+ }
469
+ require .Equal (t ,daemon1 .ID ,daemons [0 ].ProvisionerDaemon .ID )
470
+ require .Equal (t ,daemon2 .ID ,daemons [1 ].ProvisionerDaemon .ID )
471
+ require .Equal (t ,database .ProvisionerDaemonStatusOffline ,daemons [0 ].Status )
472
+ require .Equal (t ,database .ProvisionerDaemonStatusIdle ,daemons [1 ].Status )
473
+ })
474
+ }
475
+
356
476
func TestGetWorkspaceAgentUsageStats (t * testing.T ) {
357
477
t .Parallel ()
358
478