diff --git a/internal/hub/agent_connect_test.go b/internal/hub/agent_connect_test.go index 1fa87b2d..0266c7ed 100644 --- a/internal/hub/agent_connect_test.go +++ b/internal/hub/agent_connect_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "runtime" "strings" "testing" "time" @@ -34,6 +35,26 @@ func createTestHub(t testing.TB) (*Hub, *pbtests.TestApp, error) { return NewHub(testApp), 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) + } + if testApp != nil { + testApp.Cleanup() + } +} + // Helper function to create a test record func createTestRecord(app core.App, collection string, data map[string]any) (*core.Record, error) { col, err := app.FindCachedCollectionByNameOrId(collection) @@ -63,7 +84,7 @@ func TestValidateAgentHeaders(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) testCases := []struct { name string @@ -144,7 +165,7 @@ func TestGetAllFingerprintRecordsByToken(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // create test user userRecord, err := createTestUser(testApp) @@ -234,7 +255,7 @@ func TestSetFingerprint(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Create test user userRecord, err := createTestUser(testApp) @@ -314,7 +335,7 @@ func TestCreateSystemFromAgentData(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Create test user userRecord, err := createTestUser(testApp) @@ -424,7 +445,7 @@ func TestUniversalTokenFlow(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(nil, testApp) // Create test user userRecord, err := createTestUser(testApp) @@ -492,7 +513,7 @@ func TestAgentConnect(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Create test user userRecord, err := createTestUser(testApp) @@ -651,7 +672,7 @@ func TestHandleAgentConnect(t *testing.T) { if err != nil { t.Fatal(err) } - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Create test user userRecord, err := createTestUser(testApp) @@ -736,7 +757,7 @@ func TestAgentWebSocketIntegration(t *testing.T) { // Create hub and test app hub, testApp, err := createTestHub(t) require.NoError(t, err) - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Get the hub's SSH key hubSigner, err := hub.GetSSHKey("") @@ -941,6 +962,8 @@ func TestAgentWebSocketIntegration(t *testing.T) { } } + time.Sleep(20 * time.Millisecond) + // Verify fingerprint state by re-reading the specific record updatedFingerprintRecord, err := testApp.FindRecordById("fingerprints", fingerprintRecord.Id) require.NoError(t, err) @@ -975,7 +998,7 @@ func TestMultipleSystemsWithSameUniversalToken(t *testing.T) { // Create hub and test app hub, testApp, err := createTestHub(t) require.NoError(t, err) - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Get the hub's SSH key hubSigner, err := hub.GetSSHKey("") @@ -1143,6 +1166,8 @@ func TestMultipleSystemsWithSameUniversalToken(t *testing.T) { assert.Equal(t, systemCount, systemsAfterCount, "Total system count should remain the same") } + time.Sleep(20 * time.Millisecond) + // Verify that a fingerprint record exists for this fingerprint fingerprints, err := testApp.FindRecordsByFilter("fingerprints", "token = {:token} && fingerprint = {:fingerprint}", "", -1, 0, map[string]any{ "token": universalToken, @@ -1175,7 +1200,7 @@ func TestPermanentUniversalTokenFromDB(t *testing.T) { // Create hub and test app hub, testApp, err := createTestHub(t) require.NoError(t, err) - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Get the hub's SSH key hubSigner, err := hub.GetSSHKey("") @@ -1272,7 +1297,7 @@ verify: func TestFindOrCreateSystemForToken(t *testing.T) { hub, testApp, err := createTestHub(t) require.NoError(t, err) - defer testApp.Cleanup() + defer cleanupTestHub(hub, testApp) // Create test user userRecord, err := createTestUser(testApp)