fix(tests): stop system updaters on app termination

This commit is contained in:
henrygd
2026-08-17 12:08:21 -04:00
parent ae037b278e
commit 5f383c0eb1
4 changed files with 44 additions and 7 deletions

View File

@@ -56,7 +56,7 @@ func (sm *SystemManager) NewSystem(systemId string) *System {
Id: systemId, Id: systemId,
data: &system.CombinedData{}, data: &system.CombinedData{},
} }
system.ctx, system.cancel = system.getContext() system.ctx, system.cancel = system.getContext(sm.ctx)
return system return system
} }
@@ -79,7 +79,10 @@ func (sys *System) StartUpdater() {
} else { } else {
// if the system does not have a websocket connection, wait before updating // if the system does not have a websocket connection, wait before updating
// to allow the agent to connect via websocket (makes sure fingerprint is set). // to allow the agent to connect via websocket (makes sure fingerprint is set).
time.Sleep(11 * time.Second) if !waitForContext(sys.ctx, 11*time.Second) {
return
}
} }
// update immediately if system is not paused (only for ws connections) // update immediately if system is not paused (only for ws connections)
@@ -402,9 +405,9 @@ func (sys *System) setDown(originalError error) error {
return sys.manager.hub.SaveNoValidate(record) return sys.manager.hub.SaveNoValidate(record)
} }
func (sys *System) getContext() (context.Context, context.CancelFunc) { func (sys *System) getContext(ctx context.Context) (context.Context, context.CancelFunc) {
if sys.ctx == nil { if sys.ctx == nil {
sys.ctx, sys.cancel = context.WithCancel(context.Background()) sys.ctx, sys.cancel = context.WithCancel(ctx)
} }
return sys.ctx, sys.cancel return sys.ctx, sys.cancel
} }

View File

@@ -1,6 +1,7 @@
package systems package systems
import ( import (
"context"
"errors" "errors"
"fmt" "fmt"
"time" "time"
@@ -45,6 +46,8 @@ type SystemManager struct {
systems *store.Store[string, *System] // Thread-safe store of active systems systems *store.Store[string, *System] // Thread-safe store of active systems
sshConfig *ssh.ClientConfig // SSH client configuration for system connections sshConfig *ssh.ClientConfig // SSH client configuration for system connections
smartFetchMap *expirymap.ExpiryMap[smartFetchState] // Stores last SMART fetch time/result; TTL is only for cleanup smartFetchMap *expirymap.ExpiryMap[smartFetchState] // Stores last SMART fetch time/result; TTL is only for cleanup
ctx context.Context // Cancelled when the app terminates
cancel context.CancelFunc // Cancels ctx and all child system contexts
} }
// hubLike defines the interface requirements for the hub dependency. // hubLike defines the interface requirements for the hub dependency.
@@ -60,11 +63,13 @@ type hubLike interface {
// NewSystemManager creates a new SystemManager instance with the provided hub. // NewSystemManager creates a new SystemManager instance with the provided hub.
// The hub must implement the hubLike interface to provide database and alert functionality. // The hub must implement the hubLike interface to provide database and alert functionality.
func NewSystemManager(hub hubLike) *SystemManager { func NewSystemManager(hub hubLike) *SystemManager {
return &SystemManager{ sm := &SystemManager{
systems: store.New(map[string]*System{}), systems: store.New(map[string]*System{}),
hub: hub, hub: hub,
smartFetchMap: expirymap.New[smartFetchState](time.Hour), smartFetchMap: expirymap.New[smartFetchState](time.Hour),
} }
sm.ctx, sm.cancel = context.WithCancel(context.Background())
return sm
} }
// GetSystem returns a system by ID from the store // GetSystem returns a system by ID from the store
@@ -103,7 +108,9 @@ func (sm *SystemManager) Initialize() error {
sleepTime := time.Duration(delta) * time.Millisecond sleepTime := time.Duration(delta) * time.Millisecond
for _, system := range systems { for _, system := range systems {
time.Sleep(sleepTime) if !waitForContext(sm.ctx, sleepTime) {
return
}
_ = sm.AddSystem(system) _ = sm.AddSystem(system)
} }
}() }()
@@ -121,6 +128,13 @@ func (sm *SystemManager) bindEventHooks() {
sm.hub.OnRecordAfterUpdateSuccess("fingerprints").BindFunc(sm.onTokenRotated) sm.hub.OnRecordAfterUpdateSuccess("fingerprints").BindFunc(sm.onTokenRotated)
sm.hub.OnRealtimeSubscribeRequest().BindFunc(sm.onRealtimeSubscribeRequest) sm.hub.OnRealtimeSubscribeRequest().BindFunc(sm.onRealtimeSubscribeRequest)
sm.hub.OnRealtimeConnectRequest().BindFunc(sm.onRealtimeConnectRequest) sm.hub.OnRealtimeConnectRequest().BindFunc(sm.onRealtimeConnectRequest)
sm.hub.OnTerminate().BindFunc(sm.onTerminate)
}
// onTerminate cancels SystemManager context on app shutdown
func (sm *SystemManager) onTerminate(e *core.TerminateEvent) error {
sm.cancel()
return e.Next()
} }
// onTokenRotated handles fingerprint token rotation events. // onTokenRotated handles fingerprint token rotation events.
@@ -247,7 +261,7 @@ func (sm *SystemManager) AddSystem(sys *System) error {
// Initialize system for monitoring // Initialize system for monitoring
sys.manager = sm sys.manager = sm
sys.ctx, sys.cancel = sys.getContext() sys.ctx, sys.cancel = sys.getContext(sm.ctx)
sys.data = &system.CombinedData{} sys.data = &system.CombinedData{}
sm.systems.Set(sys.Id, sys) sm.systems.Set(sys.Id, sys)
@@ -372,3 +386,15 @@ func deactivateAlerts(app core.App, systemID string) error {
} }
return nil return nil
} }
// waitForContext waits for delay or returns early when ctx is cancelled.
func waitForContext(ctx context.Context, delay time.Duration) bool {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return false
case <-timer.C:
return true
}
}

View File

@@ -30,6 +30,7 @@ func TestSystemManagerNew(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
synctest.Test(t, func(t *testing.T) { synctest.Test(t, func(t *testing.T) {
sm.ResetContextForTesting()
sm.Initialize() sm.Initialize()
record, err := tests.CreateRecord(hub, "systems", map[string]any{ record, err := tests.CreateRecord(hub, "systems", map[string]any{
@@ -112,6 +113,8 @@ func TestSystemManagerNew(t *testing.T) {
assert.False(t, sm.HasSystem(record.Id), "System should not exist in the store after deletion") assert.False(t, sm.HasSystem(record.Id), "System should not exist in the store after deletion")
}) })
// The following subtests run outside the synctest bubble.
sm.ResetContextForTesting()
testOld(t, hub) testOld(t, hub)
synctest.Test(t, func(t *testing.T) { synctest.Test(t, func(t *testing.T) {

View File

@@ -117,6 +117,11 @@ func (sm *SystemManager) RemoveAllSystems() {
sm.smartFetchMap.StopCleaner() sm.smartFetchMap.StopCleaner()
} }
// ResetContextForTesting replaces the manager context for a new synctest bubble.
func (sm *SystemManager) ResetContextForTesting() {
sm.ctx, sm.cancel = context.WithCancel(context.Background())
}
func (s *System) StopUpdater() { func (s *System) StopUpdater() {
s.cancel() s.cancel()
} }