From af059350ed7b2fcaa217de0482798fd781c19715 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 5 Jul 2024 11:00:04 +0300
Subject: [PATCH] Add unit tests

---
 management/server/updatechannel.go      |   2 +-
 management/server/updatechannel_test.go | 116 ++++++++++++++++++++++++
 2 files changed, 117 insertions(+), 1 deletion(-)

diff --git a/management/server/updatechannel.go b/management/server/updatechannel.go
index 04e25d8ca69..07b5481deb9 100644
--- a/management/server/updatechannel.go
+++ b/management/server/updatechannel.go
@@ -50,7 +50,7 @@ func (p *PeersUpdateManager) SendUpdate(ctx context.Context, peerID string, upda
 
 	// skip sending sync update to the peer if there is no change in update message,
 	// it will not check on turn credential refresh as we do not send network map or client posture checks
-	if update.Update.NetworkMap != nil {
+	if update.NetworkMap != nil {
 		updated := p.handlePeerMessageUpdate(ctx, peerID, update)
 		if !updated {
 			return
diff --git a/management/server/updatechannel_test.go b/management/server/updatechannel_test.go
index 69f5b895c45..c1c250170f9 100644
--- a/management/server/updatechannel_test.go
+++ b/management/server/updatechannel_test.go
@@ -6,6 +6,8 @@ import (
 	"time"
 
 	"github.com/netbirdio/netbird/management/proto"
+	"github.com/netbirdio/netbird/management/server/posture"
+	"github.com/stretchr/testify/assert"
 )
 
 // var peersUpdater *PeersUpdateManager
@@ -77,3 +79,117 @@ func TestCloseChannel(t *testing.T) {
 		t.Error("Error closing the channel")
 	}
 }
+
+func TestHandlePeerMessageUpdate(t *testing.T) {
+	tests := []struct {
+		name               string
+		peerID             string
+		existingUpdate     *UpdateMessage
+		newUpdate          *UpdateMessage
+		expectedResult     bool
+		expectedUpdateSave bool
+	}{
+		{
+			name:   "Update message with turn credentials update",
+			peerID: "peer",
+			newUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					WiretrusteeConfig: &proto.WiretrusteeConfig{},
+				},
+			},
+			expectedResult:     true,
+			expectedUpdateSave: false,
+		},
+		{
+			name:   "Update message for peer without existing update",
+			peerID: "peer1",
+			newUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 1},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
+			},
+			expectedResult:     true,
+			expectedUpdateSave: true,
+		},
+		{
+			name:   "update message with no changes in update",
+			peerID: "peer2",
+			existingUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 1},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
+				Checks:     []*posture.Checks{},
+			},
+			newUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 1},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
+				Checks:     []*posture.Checks{},
+			},
+			expectedResult:     false,
+			expectedUpdateSave: false,
+		},
+		{
+			name:   "Update message with changes in checks",
+			peerID: "peer3",
+			existingUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 1},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
+				Checks:     []*posture.Checks{},
+			},
+			newUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 1},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
+				Checks:     []*posture.Checks{{ID: "check1"}},
+			},
+			expectedResult:     true,
+			expectedUpdateSave: true,
+		},
+		{
+			name:   "Update message with lower serial number",
+			peerID: "peer4",
+			existingUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 2},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 2}},
+			},
+			newUpdate: &UpdateMessage{
+				Update: &proto.SyncResponse{
+					NetworkMap: &proto.NetworkMap{Serial: 1},
+				},
+				NetworkMap: &NetworkMap{Network: &Network{Serial: 1}},
+			},
+			expectedResult:     false,
+			expectedUpdateSave: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := NewPeersUpdateManager(nil)
+			ctx := context.Background()
+
+			if tt.existingUpdate != nil {
+				p.peerUpdateMessage[tt.peerID] = tt.existingUpdate
+			}
+
+			result := p.handlePeerMessageUpdate(ctx, tt.peerID, tt.newUpdate)
+
+			assert.Equal(t, tt.expectedResult, result)
+
+			if tt.expectedUpdateSave {
+				assert.Equal(t, tt.newUpdate, p.peerUpdateMessage[tt.peerID])
+			} else if tt.existingUpdate != nil {
+				assert.Equal(t, tt.existingUpdate, p.peerUpdateMessage[tt.peerID])
+			}
+		})
+	}
+}