From debe65e01ef193dfb2400ae61440d64ffd7aa8b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Mon, 2 Feb 2026 11:43:35 +0100
Subject: [PATCH] Enhance unit test to verify peer count synchronization with
 debouncing and timeout handling

---
 management/server/management_test.go | 59 +++++++++++++++++++++++++---
 1 file changed, 53 insertions(+), 6 deletions(-)

diff --git a/management/server/management_test.go b/management/server/management_test.go
index 0864baadf..de02855bf 100644
--- a/management/server/management_test.go
+++ b/management/server/management_test.go
@@ -610,6 +610,7 @@ func TestSync10PeersGetUpdates(t *testing.T) {
 
 	initialPeers := 10
 	additionalPeers := 10
+	expectedPeerCount := initialPeers + additionalPeers - 1 // -1 because peer doesn't see itself
 
 	var peers []wgtypes.Key
 	for i := 0; i < initialPeers; i++ {
@@ -618,8 +619,19 @@ func TestSync10PeersGetUpdates(t *testing.T) {
 		peers = append(peers, key)
 	}
 
+	// Track the maximum peer count each peer has seen
+	type peerState struct {
+		mu           sync.Mutex
+		maxPeerCount int
+		done         bool
+	}
+	peerStates := make(map[string]*peerState)
+	for _, pk := range peers {
+		peerStates[pk.PublicKey().String()] = &peerState{}
+	}
+
 	var wg sync.WaitGroup
-	wg.Add(initialPeers + initialPeers*additionalPeers)
+	wg.Add(initialPeers) // One completion per initial peer
 
 	var syncClients []mgmtProto.ManagementService_SyncClient
 	for _, pk := range peers {
@@ -643,6 +655,9 @@ func TestSync10PeersGetUpdates(t *testing.T) {
 		syncClients = append(syncClients, s)
 
 		go func(pk wgtypes.Key, syncStream mgmtProto.ManagementService_SyncClient) {
+			pubKey := pk.PublicKey().String()
+			state := peerStates[pubKey]
+
 			for {
 				encMsg := &mgmtProto.EncryptedMessage{}
 				err := syncStream.RecvMsg(encMsg)
@@ -651,19 +666,28 @@ func TestSync10PeersGetUpdates(t *testing.T) {
 				}
 				decryptedBytes, decErr := encryption.Decrypt(encMsg.Body, ts.serverPubKey, pk)
 				if decErr != nil {
-					t.Errorf("failed to decrypt SyncResponse for peer %s: %v", pk.PublicKey().String(), decErr)
+					t.Errorf("failed to decrypt SyncResponse for peer %s: %v", pubKey, decErr)
 					return
 				}
 				resp := &mgmtProto.SyncResponse{}
 				umErr := pb.Unmarshal(decryptedBytes, resp)
 				if umErr != nil {
-					t.Errorf("failed to unmarshal SyncResponse for peer %s: %v", pk.PublicKey().String(), umErr)
+					t.Errorf("failed to unmarshal SyncResponse for peer %s: %v", pubKey, umErr)
 					return
 				}
-				// We only count if there's a new peer update
-				if len(resp.GetRemotePeers()) > 0 {
+
+				// Track the maximum peer count seen (due to debouncing, updates are coalesced)
+				peerCount := len(resp.GetRemotePeers())
+				state.mu.Lock()
+				if peerCount > state.maxPeerCount {
+					state.maxPeerCount = peerCount
+				}
+				// Signal completion when this peer has seen all expected peers
+				if !state.done && state.maxPeerCount >= expectedPeerCount {
+					state.done = true
 					wg.Done()
 				}
+				state.mu.Unlock()
 			}
 		}(pk, s)
 	}
@@ -677,7 +701,30 @@ func TestSync10PeersGetUpdates(t *testing.T) {
 		time.Sleep(time.Duration(n) * time.Millisecond)
 	}
 
-	wg.Wait()
+	// Wait for debouncer to flush final updates (debounce interval is 1000ms)
+	time.Sleep(1500 * time.Millisecond)
+
+	// Wait with timeout
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// Success - all peers received expected peer count
+	case <-time.After(5 * time.Second):
+		// Timeout - report which peers didn't receive all updates
+		t.Error("Timeout waiting for all peers to receive updates")
+		for pubKey, state := range peerStates {
+			state.mu.Lock()
+			if state.maxPeerCount < expectedPeerCount {
+				t.Errorf("Peer %s only saw %d peers, expected %d", pubKey, state.maxPeerCount, expectedPeerCount)
+			}
+			state.mu.Unlock()
+		}
+	}
 
 	for _, sc := range syncClients {
 		err := sc.CloseSend()