diff --git a/internal/alerts/alerts_battery_test.go b/internal/alerts/alerts_battery_test.go index 5361d1513..acc5a4c9e 100644 --- a/internal/alerts/alerts_battery_test.go +++ b/internal/alerts/alerts_battery_test.go @@ -5,6 +5,7 @@ package alerts_test import ( "encoding/json" "testing" + "testing/synctest" "time" "github.com/henrygd/beszel/internal/entities/system" @@ -23,7 +24,7 @@ func TestBatteryAlertLogic(t *testing.T) { defer hub.Cleanup() // Create a system - systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") require.NoError(t, err) systemRecord := systems[0] @@ -68,13 +69,12 @@ func TestBatteryAlertLogic(t *testing.T) { // Simulate system update time systemRecord.Set("updated", time.Now().UTC()) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) // Handle system alerts with high battery am := hub.GetAlertManager() - err = am.HandleSystemAlerts(systemRecord, combinedDataHigh) - require.NoError(t, err) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataHigh)) + }) // Verify alert is still NOT triggered (battery 50% is above threshold 20%) batteryAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": batteryAlert.Id}) @@ -108,15 +108,11 @@ func TestBatteryAlertLogic(t *testing.T) { // Update system timestamp systemRecord.Set("updated", time.Now().UTC()) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) // Handle system alerts with low battery - err = am.HandleSystemAlerts(systemRecord, combinedDataLow) - require.NoError(t, err) - - // Wait for the alert to be processed - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataLow)) + }) // Verify alert IS triggered (battery 15% is below threshold 20%) batteryAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": batteryAlert.Id}) @@ -150,15 +146,11 @@ func TestBatteryAlertLogic(t *testing.T) { // Update system timestamp systemRecord.Set("updated", time.Now().UTC()) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) // Handle system alerts with recovered battery - err = am.HandleSystemAlerts(systemRecord, combinedDataRecovered) - require.NoError(t, err) - - // Wait for the alert to be processed - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataRecovered)) + }) // Verify alert is now resolved (battery 25% is above threshold 20%) batteryAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": batteryAlert.Id}) @@ -172,7 +164,7 @@ func TestBatteryAlertNoBattery(t *testing.T) { defer hub.Cleanup() // Create a system - systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") require.NoError(t, err) systemRecord := systems[0] @@ -206,16 +198,12 @@ func TestBatteryAlertNoBattery(t *testing.T) { // Simulate system update time systemRecord.Set("updated", time.Now().UTC()) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) // Handle system alerts am := hub.GetAlertManager() - err = am.HandleSystemAlerts(systemRecord, combinedData) - require.NoError(t, err) - - // Wait a moment for processing - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedData)) + }) // Verify alert is NOT triggered (no battery data should skip the alert) batteryAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": batteryAlert.Id}) @@ -230,7 +218,7 @@ func TestBatteryAlertAveragedSamples(t *testing.T) { defer hub.Cleanup() // Create a system - systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") require.NoError(t, err) systemRecord := systems[0] @@ -302,15 +290,11 @@ func TestBatteryAlertAveragedSamples(t *testing.T) { // Update system timestamp systemRecord.Set("updated", now) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) // Handle system alerts - should trigger because average battery is below threshold - err = am.HandleSystemAlerts(systemRecord, combinedDataLow) - require.NoError(t, err) - - // Wait for alert processing - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataLow)) + }) // Verify alert IS triggered (average battery 15% is below threshold 25%) batteryAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": batteryAlert.Id}) @@ -368,15 +352,11 @@ func TestBatteryAlertAveragedSamples(t *testing.T) { // Update system timestamp to the new time window systemRecord.Set("updated", newNow) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) // Handle system alerts - should resolve because average battery is now above threshold - err = am.HandleSystemAlerts(systemRecord, combinedDataHigh) - require.NoError(t, err) - - // Wait for alert processing - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataHigh)) + }) // Verify alert is resolved (average battery 50% is above threshold 25%) batteryAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": batteryAlert.Id}) diff --git a/internal/alerts/alerts_disk_test.go b/internal/alerts/alerts_disk_test.go index 72f094d03..5c9a57389 100644 --- a/internal/alerts/alerts_disk_test.go +++ b/internal/alerts/alerts_disk_test.go @@ -5,6 +5,7 @@ package alerts_test import ( "encoding/json" "testing" + "testing/synctest" "time" "github.com/henrygd/beszel/internal/entities/system" @@ -22,7 +23,7 @@ func TestDiskAlertExtraFsMultiMinute(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) defer hub.Cleanup() - systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") require.NoError(t, err) systemRecord := systems[0] @@ -83,13 +84,10 @@ func TestDiskAlertExtraFsMultiMinute(t *testing.T) { } systemRecord.Set("updated", now) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) - err = am.HandleSystemAlerts(systemRecord, combinedDataHigh) - require.NoError(t, err) - - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataHigh)) + }) diskAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": diskAlert.Id}) require.NoError(t, err) @@ -140,13 +138,10 @@ func TestDiskAlertExtraFsMultiMinute(t *testing.T) { } systemRecord.Set("updated", newNow) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) - err = am.HandleSystemAlerts(systemRecord, combinedDataLow) - require.NoError(t, err) - - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataLow)) + }) diskAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": diskAlert.Id}) require.NoError(t, err) diff --git a/internal/alerts/alerts_quiet_hours_test.go b/internal/alerts/alerts_quiet_hours_test.go index 328e99cea..b58043462 100644 --- a/internal/alerts/alerts_quiet_hours_test.go +++ b/internal/alerts/alerts_quiet_hours_test.go @@ -323,9 +323,14 @@ func TestAlertSilencedMultiUser(t *testing.T) { func TestAlertSilencedWithActualAlert(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) // Create a system systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") diff --git a/internal/alerts/alerts_shutdown_test.go b/internal/alerts/alerts_shutdown_test.go new file mode 100644 index 000000000..5ca3f8b1d --- /dev/null +++ b/internal/alerts/alerts_shutdown_test.go @@ -0,0 +1,48 @@ +//go:build testing + +package alerts_test + +import ( + "sync" + "testing" + + "github.com/henrygd/beszel/internal/alerts" + beszelTests "github.com/henrygd/beszel/internal/tests" + "github.com/stretchr/testify/require" +) + +func TestStandaloneAlertManagerStopsBeforeDatabaseCleanup(t *testing.T) { + for _, state := range []string{"pending", "already stopped", "delivered"} { + t.Run(state, func(t *testing.T) { + hub, user := beszelTests.GetHubWithUser(t) + cleanup := sync.OnceFunc(hub.Cleanup) + defer cleanup() + setStatusAlertEmail(t, hub, user.Id, "shutdown@example.com") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") + require.NoError(t, err) + _, err = beszelTests.CreateRecord(hub, "alerts", map[string]any{ + "name": "Status", "system": systems[0].Id, "user": user.Id, "min": 1, + }) + require.NoError(t, err) + + am := alerts.NewTestAlertManagerWithoutWorker(hub) + defer am.Stop() + require.NoError(t, am.HandleStatusAlerts("down", systems[0])) + require.Equal(t, 1, am.GetPendingAlertsCount()) + switch state { + case "already stopped": + am.Stop() + case "delivered": + am.ForceExpirePendingAlerts() + processed, err := am.ProcessPendingAlerts() + require.NoError(t, err) + require.Len(t, processed, 1) + require.Equal(t, 1, hub.TestMailer.TotalSend()) + } + + cleanup() + // No timer may retain the disposed app until its eventual deadline. + require.Zero(t, am.GetPendingAlertsCount()) + }) + } +} diff --git a/internal/alerts/alerts_status_test.go b/internal/alerts/alerts_status_test.go index ecfc0c0fb..a4cddd26b 100644 --- a/internal/alerts/alerts_status_test.go +++ b/internal/alerts/alerts_status_test.go @@ -30,9 +30,14 @@ func setStatusAlertEmail(t *testing.T, hub core.App, userID, email string) { func TestStatusAlerts(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) systems, err := beszelTests.CreateSystems(hub, 4, user.Id, "paused") assert.NoError(t, err) @@ -236,9 +241,14 @@ func TestHandleStatusAlertsDoesNotSendRecoveryWhileDownIsOnlyPending(t *testing. func TestStatusAlertTimerCancellationPreventsBoundaryDelivery(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) userSettings, err := hub.FindFirstRecordByFilter("user_settings", "user={:user}", map[string]any{"user": user.Id}) require.NoError(t, err) @@ -339,9 +349,14 @@ func TestStatusAlertDownFiresAfterDelayExpires(t *testing.T) { func TestStatusAlertMultipleUsersRespectDifferentMinutes(t *testing.T) { hub, user1 := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) setStatusAlertEmail(t, hub, user1.Id, "user1@example.com") @@ -427,9 +442,14 @@ func TestStatusAlertMultipleUsersRespectDifferentMinutes(t *testing.T) { func TestStatusAlertMultipleUsersRecoveryBetweenMinutesOnlyAlertsEarlierUser(t *testing.T) { hub, user1 := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) setStatusAlertEmail(t, hub, user1.Id, "user1@example.com") @@ -821,9 +841,14 @@ func TestResolveStatusAlerts(t *testing.T) { func TestAlertsHistoryStatus(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) // Create a system systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") @@ -888,9 +913,14 @@ func TestAlertsHistoryStatus(t *testing.T) { func TestStatusAlertClearedBeforeSend(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) // Create a system systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") diff --git a/internal/alerts/alerts_test.go b/internal/alerts/alerts_test.go index a40bf9136..7945fd146 100644 --- a/internal/alerts/alerts_test.go +++ b/internal/alerts/alerts_test.go @@ -16,9 +16,14 @@ import ( func TestAlertsHistory(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) + defer hub.Cleanup() synctest.Test(t, func(t *testing.T) { - defer hub.Cleanup() + t.Cleanup(func() { + synctest.Wait() + hub.GetAlertManager().Stop() + }) + hub.GetSystemManager().ResetContextForTesting(t) // Create systems and alerts systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") diff --git a/internal/alerts/alerts_test_helpers.go b/internal/alerts/alerts_test_helpers.go index 89f067c2a..7fa7dd44c 100644 --- a/internal/alerts/alerts_test_helpers.go +++ b/internal/alerts/alerts_test_helpers.go @@ -10,11 +10,17 @@ import ( ) func NewTestAlertManagerWithoutWorker(app hubLike) *AlertManager { - return &AlertManager{ + am := &AlertManager{ hub: app, alertsCache: NewAlertsCache(app), networkMonitors: newNetworkMonitorCache(app), } + // Standalone managers can own status timers even without the serve hooks. + app.OnTerminate().BindFunc(func(e *core.TerminateEvent) error { + am.Stop() + return e.Next() + }) + return am } // GetSystemAlertsCache returns the internal system alerts cache. diff --git a/internal/alerts/alerts_zfs_disk_test.go b/internal/alerts/alerts_zfs_disk_test.go index a248a89cc..46ad528b1 100644 --- a/internal/alerts/alerts_zfs_disk_test.go +++ b/internal/alerts/alerts_zfs_disk_test.go @@ -5,6 +5,7 @@ package alerts_test import ( "encoding/json" "testing" + "testing/synctest" "time" "github.com/henrygd/beszel/internal/entities/system" @@ -23,7 +24,7 @@ func TestDiskAlertZfsPoolMultiMinute(t *testing.T) { hub, user := beszelTests.GetHubWithUser(t) defer hub.Cleanup() - systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") require.NoError(t, err) systemRecord := systems[0] @@ -80,13 +81,10 @@ func TestDiskAlertZfsPoolMultiMinute(t *testing.T) { } systemRecord.Set("updated", now) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) - err = am.HandleSystemAlerts(systemRecord, combinedDataHigh) - require.NoError(t, err) - - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataHigh)) + }) diskAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": diskAlert.Id}) require.NoError(t, err) @@ -130,13 +128,10 @@ func TestDiskAlertZfsPoolMultiMinute(t *testing.T) { } systemRecord.Set("updated", newNow) - err = hub.SaveNoValidate(systemRecord) - require.NoError(t, err) - err = am.HandleSystemAlerts(systemRecord, combinedDataLow) - require.NoError(t, err) - - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, am.HandleSystemAlerts(systemRecord, combinedDataLow)) + }) diskAlert, err = hub.FindFirstRecordByFilter("alerts", "id={:id}", dbx.Params{"id": diskAlert.Id}) require.NoError(t, err) @@ -147,7 +142,7 @@ func TestDiskAlertZfsPoolMultiMinute(t *testing.T) { func TestDiskAlertIgnoresRawPool(t *testing.T) { for _, minutes := range []int{0, 2} { hub, user := beszelTests.GetHubWithUser(t) - systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "up") + systems, err := beszelTests.CreateSystems(hub, 1, user.Id, "paused") require.NoError(t, err) alert, err := beszelTests.CreateRecord(hub, "alerts", map[string]any{"name": "Disk", "system": systems[0].Id, "user": user.Id, "value": 80, "min": minutes}) require.NoError(t, err) @@ -160,16 +155,18 @@ func TestDiskAlertIgnoresRawPool(t *testing.T) { record.SetRaw("created", time.Now().UTC().Add(offset*time.Second).Format(types.DefaultDateLayout)) require.NoError(t, hub.SaveNoValidate(record)) } - require.NoError(t, hub.GetAlertManager().HandleSystemAlerts(systems[0], &system.CombinedData{Stats: system.Stats{ZfsPools: pools}})) - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, hub.GetAlertManager().HandleSystemAlerts(systems[0], &system.CombinedData{Stats: system.Stats{ZfsPools: pools}})) + }) record, err := hub.FindRecordById("alerts", alert.Id) require.NoError(t, err) assert.False(t, record.GetBool("triggered")) if minutes > 0 { // A current usable sample must not make raw historical values eligible. pools["btrfs"].Raw = false - require.NoError(t, hub.GetAlertManager().HandleSystemAlerts(systems[0], &system.CombinedData{Stats: system.Stats{ZfsPools: pools}})) - time.Sleep(20 * time.Millisecond) + synctest.Test(t, func(t *testing.T) { + require.NoError(t, hub.GetAlertManager().HandleSystemAlerts(systems[0], &system.CombinedData{Stats: system.Stats{ZfsPools: pools}})) + }) record, err = hub.FindRecordById("alerts", alert.Id) require.NoError(t, err) assert.False(t, record.GetBool("triggered")) diff --git a/internal/hub/agent_connect_test.go b/internal/hub/agent_connect_test.go index 4bbc9767d..d4ec4e8b0 100644 --- a/internal/hub/agent_connect_test.go +++ b/internal/hub/agent_connect_test.go @@ -9,7 +9,6 @@ import ( "net/http/httptest" "os" "path/filepath" - "runtime" "strings" "testing" "time" @@ -32,23 +31,18 @@ func createTestHub(t testing.TB) (*Hub, *pbtests.TestApp, error) { if err != nil { return nil, nil, err } - return NewHub(testApp), testApp, err + hub := NewHub(testApp) + if err := hub.sm.InitializeSSHConfigForTesting(); err != nil { + cleanupTestHub(hub, testApp) + return nil, nil, err + } + return hub, testApp, nil } // cleanupTestHub stops background system goroutines before tearing down the app. func cleanupTestHub(hub *Hub, testApp *pbtests.TestApp) { if hub != nil { - sm := hub.GetSystemManager() - sm.RemoveAllSystems() - // Give updater goroutines a brief window to observe cancellation before DB teardown. - for range 20 { - if sm.GetSystemCount() == 0 { - break - } - runtime.Gosched() - time.Sleep(5 * time.Millisecond) - } - time.Sleep(20 * time.Millisecond) + hub.GetSystemManager().RemoveAllSystems() } if testApp != nil { testApp.Cleanup() diff --git a/internal/hub/systems/network_monitor_sync_test.go b/internal/hub/systems/network_monitor_sync_test.go index ea50d0979..4e0c26656 100644 --- a/internal/hub/systems/network_monitor_sync_test.go +++ b/internal/hub/systems/network_monitor_sync_test.go @@ -91,6 +91,7 @@ func TestNetworkMonitorReconnectSync(t *testing.T) { require.NoError(t, app.SaveNoValidate(probe)) sm := NewSystemManager(stubHub{app}) + require.NoError(t, sm.createSSHClientConfig()) t.Cleanup(func() { sm.cancel() _ = sm.RemoveSystem(sys.Id) diff --git a/internal/hub/systems/system_manager.go b/internal/hub/systems/system_manager.go index 6c14b8d9a..7bd27ab71 100644 --- a/internal/hub/systems/system_manager.go +++ b/internal/hub/systems/system_manager.go @@ -45,6 +45,7 @@ var errSystemExists = errors.New("system exists") // SystemManager manages a collection of monitored systems and their connections. // It handles system lifecycle, status updates, and maintains both SSH and WebSocket connections. type SystemManager struct { + updaters sync.WaitGroup // Tracks updater completion independently of store membership. hub hubLike // Hub interface for database and alert operations systems *store.Store[string, *System] // Thread-safe store of active systems sshConfig *ssh.ClientConfig // SSH client configuration for system connections @@ -293,7 +294,7 @@ func (sm *SystemManager) AddSystem(sys *System) error { sm.systems.Set(sys.Id, sys) // Start monitoring in background - go sys.StartUpdater() + sm.updaters.Go(sys.StartUpdater) return nil } diff --git a/internal/hub/systems/systems_test.go b/internal/hub/systems/systems_test.go index ce87f3a44..d55c4cdc7 100644 --- a/internal/hub/systems/systems_test.go +++ b/internal/hub/systems/systems_test.go @@ -62,13 +62,14 @@ func TestSystemManagerNew(t *testing.T) { require.NoError(t, err) synctest.Test(t, func(t *testing.T) { - sm.ResetContextForTesting() + sm.ResetContextForTesting(t) sm.Initialize() + // Use a closed loopback endpoint so fake-clock tests do not initialise shared DNS state. record, err := tests.CreateRecord(hub, "systems", map[string]any{ "name": "it-was-coney-island", - "host": "the-playground-of-the-world", - "port": "33914", + "host": "127.0.0.1", + "port": "0", "users": []string{user.Id}, }) require.NoError(t, err) @@ -109,8 +110,8 @@ func TestSystemManagerNew(t *testing.T) { // let's also make sure a system is removed from the store when the record is deleted record, err = tests.CreateRecord(hub, "systems", map[string]any{ "name": "there-was-no-place-like-it", - "host": "in-the-whole-world", - "port": "33914", + "host": "127.0.0.1", + "port": "0", "users": []string{user.Id}, }) require.NoError(t, err) @@ -146,7 +147,7 @@ func TestSystemManagerNew(t *testing.T) { }) // The following subtests run outside the synctest bubble. - sm.ResetContextForTesting() + sm.ResetContextForTesting(t) testOld(t, hub) synctest.Test(t, func(t *testing.T) { @@ -498,7 +499,7 @@ func testOld(t *testing.T, hub *tests.TestHub) { assert.NoError(t, err) assert.NotNil(t, newCtx, "New system context should not be nil") assert.NotNil(t, newCancel, "New system cancel function should not be nil") - assert.NotEqual(t, originalCtx, newCtx, "New context should be different from original") + assert.NotSame(t, originalCtx, newCtx, "New context should be different from original") // Clean up err = sm.RemoveSystem(record.Id) diff --git a/internal/hub/systems/systems_test_helpers.go b/internal/hub/systems/systems_test_helpers.go index 9e6c48601..57aa9385f 100644 --- a/internal/hub/systems/systems_test_helpers.go +++ b/internal/hub/systems/systems_test_helpers.go @@ -5,6 +5,7 @@ package systems import ( "context" "fmt" + "testing" entities "github.com/henrygd/beszel/internal/entities/system" "github.com/henrygd/beszel/internal/entities/systemd" @@ -21,6 +22,11 @@ func backgroundSmartFetchEnabled() bool { return false } // Background ZFS fetching follows the same policy as SMART fetching. func backgroundZfsFetchEnabled() bool { return false } +// InitializeSSHConfigForTesting prepares SSH without registering hooks or starting systems. +func (sm *SystemManager) InitializeSSHConfigForTesting() error { + return sm.createSSHClientConfig() +} + // TESTING ONLY: GetSystemCount returns the number of systems in the store func (sm *SystemManager) GetSystemCount() int { return sm.systems.Length() @@ -106,25 +112,32 @@ func (sm *SystemManager) SetSystemStatusInDB(systemID string, status string) boo record.Set("status", status) err = sm.hub.Save(record) - if err != nil { - return false - } - - return true + return err == nil } -// TESTING ONLY: RemoveAllSystems removes all systems from the store +// TESTING ONLY: RemoveAllSystems cancels and joins updaters before database cleanup. +// Callers must stop producers that can add systems before calling it. func (sm *SystemManager) RemoveAllSystems() { for _, system := range sm.systems.GetAll() { sm.RemoveSystem(system.Id) } + // Removed or replaced systems may still be finishing their last update. + sm.updaters.Wait() sm.smartFetchMap.StopCleaner() sm.zfsFetchMap.StopCleaner() } // ResetContextForTesting replaces the manager context for a new synctest bubble. -func (sm *SystemManager) ResetContextForTesting() { +func (sm *SystemManager) ResetContextForTesting(t testing.TB) { + t.Helper() + sm.cancel() + sm.updaters.Wait() sm.ctx, sm.cancel = context.WithCancel(context.Background()) + cancel := sm.cancel + t.Cleanup(func() { + cancel() + sm.updaters.Wait() + }) } func (s *System) StopUpdater() { diff --git a/internal/hub/systems/updater_cleanup_test.go b/internal/hub/systems/updater_cleanup_test.go new file mode 100644 index 000000000..31914a161 --- /dev/null +++ b/internal/hub/systems/updater_cleanup_test.go @@ -0,0 +1,106 @@ +//go:build testing + +package systems + +import ( + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/blang/semver" + "github.com/henrygd/beszel/internal/entities/monitor" + esystem "github.com/henrygd/beszel/internal/entities/system" + "github.com/henrygd/beszel/internal/hub/ws" + "github.com/lxzan/gws" + "github.com/pocketbase/pocketbase/core" + "github.com/stretchr/testify/require" +) + +type blockedUpdateHub struct { + stubHub + entered chan struct{} + release chan struct{} + finished chan struct{} +} + +func (h *blockedUpdateHub) HandleNetworkMonitorAlerts(*core.Record, map[string]monitor.Result) error { + close(h.entered) + <-h.release + defer close(h.finished) + // This callback is still allowed to use the database after cancellation. + _, err := h.FindRecordsByFilter("alerts", "", "", 1, 0) + return err +} + +func TestRemoveAllSystemsWaitsForUpdater(t *testing.T) { + for _, removed := range []bool{false, true} { + name := "active" + if removed { + name = "already removed" + } + t.Run(name, func(t *testing.T) { + sys, app := newTestSystemWithHub(t) + h := &blockedUpdateHub{stubHub: stubHub{app}, entered: make(chan struct{}), release: make(chan struct{}), finished: make(chan struct{})} + release := sync.OnceFunc(func() { close(h.release) }) + // Release first if a setup assertion fails, then stop workers before the DB. + sm := NewSystemManager(h) + t.Cleanup(sm.RemoveAllSystems) + t.Cleanup(release) + + connections := make(chan *ws.WsConn, 1) + upgrader := gws.NewUpgrader(&monitorSyncServer{}, nil) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r) + if err != nil { + t.Error(err) + return + } + wsConn := ws.NewWsConnection(conn, semver.MustParse("0.20.0")) + conn.Session().Store("wsConn", wsConn) + connections <- wsConn + conn.ReadLoop() + })) + t.Cleanup(server.Close) + client := &sequenceDataClient{responses: make(chan esystem.CombinedData, 1)} + client.responses <- esystem.CombinedData{} + conn, _, err := gws.NewClient(client, &gws.ClientOption{Addr: "ws" + strings.TrimPrefix(server.URL, "http")}) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.NetConn().Close() }) + go conn.ReadLoop() + select { + case sys.WsConn = <-connections: + case <-time.After(3 * time.Second): + t.Fatal("websocket was not established") + } + sys.Host = "127.0.0.1" + require.NoError(t, sm.AddSystem(sys)) + select { + case <-h.entered: + case <-time.After(3 * time.Second): + t.Fatal("updater did not reach network alerts") + } + if removed { + require.NoError(t, sm.RemoveSystem(sys.Id)) + } + done := make(chan struct{}) + go func() { sm.RemoveAllSystems(); close(done) }() + <-sys.ctx.Done() + select { + case <-done: + t.Error("cleanup returned while an updater still needed the database") + case <-time.After(100 * time.Millisecond): + } + release() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("cleanup did not finish after the updater was released") + } + <-h.finished + require.Zero(t, sm.GetSystemCount()) + }) + } +}