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

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)

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
 }
 

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)

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() {

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())
+		})
+	}
+}