@@ -353,6 +353,126 @@ func TestGetEligibleProvisionerDaemonsByProvisionerJobIDs(t *testing.T) {
353353})
354354}
355355
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+
356476func TestGetWorkspaceAgentUsageStats (t * testing.T ) {
357477t .Parallel ()
358478