test: initialise and stop hub fixture workers before cleanup (#2463)

2026-09-29 04:47:46 +02:00 · 2026-09-29 12:42:36 +13:00
parent 130365f5d3
commit 8b38db1ba0
14 changed files with 292 additions and 110 deletions
--- 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})
--- 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)
--- 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")
--- a/internal/alerts/alerts_shutdown_test.go
+++ 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())
+		})
+	}
+}
--- 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")
--- 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")
--- 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.
--- 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"))
--- 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()
--- 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)
--- 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
 }

--- 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)
--- 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() {
--- a/internal/hub/systems/updater_cleanup_test.go
+++ 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())
+		})
+	}
+}