[management] extract peer update logic and wrap it in tests (#7338)

* extract peer update loop into a dedicated struct and wrap it in tests Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * make linter happy Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> --------- Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
2026-09-12 17:59:06 +02:00 · 2026-09-11 17:07:26 +02:00
parent 794956a7a3
commit f422c41654
9 changed files with 589 additions and 86 deletions
@@ -6,6 +6,16 @@ import (
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )

+type Encrypter interface {
+	EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error)
+}
+
+type DefaultEncrypter struct{}
+
+func (e DefaultEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) {
+	return EncryptMessage(remotePubKey, ourPrivateKey, message)
+}
+
 // EncryptMessage encrypts a body of the given protobuf Message
 func EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) {
 	byteResp, err := pb.Marshal(message)
@@ -0,0 +1,135 @@
+package grpc
+
+import (
+	"context"
+	"time"
+
+	"github.com/netbirdio/netbird/encryption"
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/management/server/telemetry"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	log "github.com/sirupsen/logrus"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func PeerUpdateHandlerFactory(
+	peerKey wgtypes.Key,
+	updates chan *network_map.UpdateMessage,
+	secretsManager SecretsManager,
+	srv proto.ManagementService_SyncServer,
+	cleanupfunc func()) *PeerUpdateHandler {
+	return &PeerUpdateHandler{
+		peerKey:        peerKey,
+		updates:        updates,
+		secretsManager: secretsManager,
+		srv:            srv,
+		encrypter:      encryption.DefaultEncrypter{},
+		debouncer:      NewUpdateDebouncer(1000 * time.Millisecond),
+		cleanupFunc:    cleanupfunc,
+	}
+}
+
+// PeerUpdateHandler sends updates to the connected peer until the updates channel is closed.
+// It implements a backpressure mechanism that sends the first update immediately,
+// then debounces subsequent rapid updates, ensuring only the latest update is sent
+// after a quiet period.
+type PeerUpdateHandler struct {
+	peerKey        wgtypes.Key
+	updates        chan *network_map.UpdateMessage
+	appMetrics     telemetry.AppMetrics
+	secretsManager SecretsManager
+	srv            syncSender
+	encrypter      encryption.Encrypter
+	debouncer      Debouncer
+	cleanupFunc    func()
+}
+
+func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler {
+	pu.appMetrics = appMetrics
+	return pu
+}
+
+//go:generate go tool mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
+type syncSender interface {
+	Send(*proto.EncryptedMessage) error
+	Context() context.Context
+}
+
+func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error {
+	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", pu.peerKey.String())
+
+	defer pu.debouncer.Stop()
+
+	for {
+		select {
+		// condition when there are some updates
+		// todo set the updates channel size to 1
+		case update, open := <-pu.updates:
+			if pu.appMetrics != nil {
+				pu.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(pu.updates) + 1)
+			}
+
+			if !open {
+				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", pu.peerKey.String())
+				pu.cleanupFunc()
+				return nil
+			}
+
+			log.WithContext(ctx).Tracef("received an update for peer %s", pu.peerKey.String())
+			if pu.debouncer.ProcessUpdate(update) {
+				// Send immediately (first update or after quiet period)
+				if err := pu.SendUpdate(ctx, update); err != nil {
+					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
+					return err
+				}
+			}
+
+		// Timer expired - quiet period reached, send pending updates if any
+		case <-pu.debouncer.TimerChannel():
+			pendingUpdates := pu.debouncer.GetPendingUpdates()
+			if len(pendingUpdates) == 0 {
+				continue
+			}
+			log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), pu.peerKey.String())
+			for _, pendingUpdate := range pendingUpdates {
+				if err := pu.SendUpdate(ctx, pendingUpdate); err != nil {
+					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
+					return err
+				}
+			}
+
+		// condition when client <-> server connection has been terminated
+		case <-pu.srv.Context().Done():
+			// happens when connection drops, e.g. client disconnects
+			log.WithContext(ctx).Debugf("stream of peer %s has been closed", pu.peerKey.String())
+			pu.cleanupFunc()
+			return pu.srv.Context().Err()
+		}
+	}
+}
+
+func (pu *PeerUpdateHandler) SendUpdate(ctx context.Context, update *network_map.UpdateMessage) error {
+	key, err := pu.secretsManager.GetWGKey()
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed processing update message")
+	}
+
+	encryptedResp, err := pu.encrypter.EncryptMessage(pu.peerKey, key, update.Update)
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed processing update message")
+	}
+	err = pu.srv.Send(&proto.EncryptedMessage{
+		WgPubKey: key.PublicKey().String(),
+		Body:     encryptedResp,
+	})
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed sending update message")
+	}
+	log.WithContext(ctx).Tracef("sent an update to peer %s", pu.peerKey.String())
+	return nil
+}
@@ -0,0 +1,155 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	pb "github.com/golang/protobuf/proto" //nolint
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+func TestSendPeerUpdates_FirstUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	msg := network_map.UpdateMessage{
+		Update: &proto.SyncResponse{Version: 1},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx := context.TODO()
+	srvKey := mustGenerateKey(t)
+	// mock a first update, should send it right away
+	updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true)
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+	syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	pu.updates <- &msg
+	close(pu.updates)
+	wg.Wait()
+}
+
+func TestSendPeerUpdates_TimerUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	msg := network_map.UpdateMessage{
+		Update: &proto.SyncResponse{Version: 1},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx := context.TODO()
+	srvKey := mustGenerateKey(t)
+	updateDebouncer.EXPECT().GetPendingUpdates().Return([]*network_map.UpdateMessage{&msg})
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+	syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	timeCh <- time.Now()
+	close(pu.updates)
+	wg.Wait()
+}
+
+func TestSendPeerUpdates_ServerContextDone(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx, cancel := context.WithCancel(context.TODO())
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	cancel()
+	wg.Wait()
+}
+
+func mustGenerateKey(t *testing.T) wgtypes.Key {
+	t.Helper()
+	k, err := wgtypes.GenerateKey()
+	assert.NoError(t, err)
+	return k
+}
+
+func mustMarshal(t *testing.T, msg *network_map.UpdateMessage) []byte {
+	t.Helper()
+	r, err := pb.Marshal(msg.Update)
+	assert.NoError(t, err)
+	return r
+}
+
+type testEncrypter struct{}
+
+func (testEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) {
+	return pb.Marshal(message)
+}
+
+type pbMatcher struct {
+	x pb.Message
+}
+
+func (pbm pbMatcher) Matches(x any) bool {
+	msg, ok := x.(pb.Message)
+	if !ok {
+		return false
+	}
+	return pb.Equal(pbm.x, msg)
+}
+
+func (pbm pbMatcher) String() string {
+	return fmt.Sprintf("is equal to %s (%T)", pbm.x, pbm.x)
+}
@@ -337,7 +337,8 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S

 	s.syncSem.Add(-1)

-	return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv, syncStart)
+	return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart) }).
+		WithMetrics(s.appMetrics).HandleUpdates(ctx)
 }

 func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementService_JobServer) (wgtypes.Key, error) {
@@ -404,91 +405,6 @@ func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgt
 	}
 }

-// handleUpdates sends updates to the connected peer until the updates channel is closed.
-// It implements a backpressure mechanism that sends the first update immediately,
-// then debounces subsequent rapid updates, ensuring only the latest update is sent
-// after a quiet period.
-func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
-	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String())
-
-	// Create a debouncer for this peer connection
-	debouncer := NewUpdateDebouncer(1000 * time.Millisecond)
-	defer debouncer.Stop()
-
-	for {
-		select {
-		// condition when there are some updates
-		// todo set the updates channel size to 1
-		case update, open := <-updates:
-			if s.appMetrics != nil {
-				s.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(updates) + 1)
-			}
-
-			if !open {
-				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String())
-				s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-				return nil
-			}
-
-			log.WithContext(ctx).Tracef("received an update for peer %s", peerKey.String())
-			if debouncer.ProcessUpdate(update) {
-				// Send immediately (first update or after quiet period)
-				if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime); err != nil {
-					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
-					return err
-				}
-			}
-
-		// Timer expired - quiet period reached, send pending updates if any
-		case <-debouncer.TimerChannel():
-			pendingUpdates := debouncer.GetPendingUpdates()
-			if len(pendingUpdates) == 0 {
-				continue
-			}
-			log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), peerKey.String())
-			for _, pendingUpdate := range pendingUpdates {
-				if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime); err != nil {
-					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
-					return err
-				}
-			}
-
-		// condition when client <-> server connection has been terminated
-		case <-srv.Context().Done():
-			// happens when connection drops, e.g. client disconnects
-			log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
-			s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-			return srv.Context().Err()
-		}
-	}
-}
-
-// sendUpdate encrypts the update message using the peer key and the server's wireguard key,
-// then sends the encrypted message to the connected peer via the sync server.
-func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update)
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-	err = srv.Send(&proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
-		Body:     encryptedResp,
-	})
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed sending update message")
-	}
-	log.WithContext(ctx).Tracef("sent an update to peer %s", peerKey.String())
-	return nil
-}
-
 // sendJob encrypts the update message using the peer key and the server's wireguard key,
 // then sends the encrypted message to the connected peer via the sync server.
 func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Event, srv proto.ManagementService_JobServer) error {
@@ -0,0 +1,70 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./peer_update_handler.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	context "context"
+	reflect "reflect"
+
+	proto "github.com/netbirdio/netbird/shared/management/proto"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MocksyncSender is a mock of syncSender interface.
+type MocksyncSender struct {
+	ctrl     *gomock.Controller
+	recorder *MocksyncSenderMockRecorder
+	isgomock struct{}
+}
+
+// MocksyncSenderMockRecorder is the mock recorder for MocksyncSender.
+type MocksyncSenderMockRecorder struct {
+	mock *MocksyncSender
+}
+
+// NewMocksyncSender creates a new mock instance.
+func NewMocksyncSender(ctrl *gomock.Controller) *MocksyncSender {
+	mock := &MocksyncSender{ctrl: ctrl}
+	mock.recorder = &MocksyncSenderMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MocksyncSender) EXPECT() *MocksyncSenderMockRecorder {
+	return m.recorder
+}
+
+// Context mocks base method.
+func (m *MocksyncSender) Context() context.Context {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Context")
+	ret0, _ := ret[0].(context.Context)
+	return ret0
+}
+
+// Context indicates an expected call of Context.
+func (mr *MocksyncSenderMockRecorder) Context() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MocksyncSender)(nil).Context))
+}
+
+// Send mocks base method.
+func (m *MocksyncSender) Send(arg0 *proto.EncryptedMessage) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Send", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Send indicates an expected call of Send.
+func (mr *MocksyncSenderMockRecorder) Send(arg0 any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MocksyncSender)(nil).Send), arg0)
+}
@@ -25,6 +25,8 @@ import (
 const defaultDuration = 12 * time.Hour

 // SecretsManager used to manage TURN and relay secrets
+//
+//go:generate go tool mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
 type SecretsManager interface {
 	GenerateTurnToken() (*Token, error)
 	GenerateRelayToken() (*Token, error)
@@ -0,0 +1,111 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./token_mgr.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	context "context"
+	reflect "reflect"
+
+	gomock "go.uber.org/mock/gomock"
+	wgtypes "golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+// MockSecretsManager is a mock of SecretsManager interface.
+type MockSecretsManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockSecretsManagerMockRecorder
+	isgomock struct{}
+}
+
+// MockSecretsManagerMockRecorder is the mock recorder for MockSecretsManager.
+type MockSecretsManagerMockRecorder struct {
+	mock *MockSecretsManager
+}
+
+// NewMockSecretsManager creates a new mock instance.
+func NewMockSecretsManager(ctrl *gomock.Controller) *MockSecretsManager {
+	mock := &MockSecretsManager{ctrl: ctrl}
+	mock.recorder = &MockSecretsManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockSecretsManager) EXPECT() *MockSecretsManagerMockRecorder {
+	return m.recorder
+}
+
+// CancelRefresh mocks base method.
+func (m *MockSecretsManager) CancelRefresh(peerKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "CancelRefresh", peerKey)
+}
+
+// CancelRefresh indicates an expected call of CancelRefresh.
+func (mr *MockSecretsManagerMockRecorder) CancelRefresh(peerKey any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRefresh", reflect.TypeOf((*MockSecretsManager)(nil).CancelRefresh), peerKey)
+}
+
+// GenerateRelayToken mocks base method.
+func (m *MockSecretsManager) GenerateRelayToken() (*Token, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GenerateRelayToken")
+	ret0, _ := ret[0].(*Token)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GenerateRelayToken indicates an expected call of GenerateRelayToken.
+func (mr *MockSecretsManagerMockRecorder) GenerateRelayToken() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateRelayToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateRelayToken))
+}
+
+// GenerateTurnToken mocks base method.
+func (m *MockSecretsManager) GenerateTurnToken() (*Token, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GenerateTurnToken")
+	ret0, _ := ret[0].(*Token)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GenerateTurnToken indicates an expected call of GenerateTurnToken.
+func (mr *MockSecretsManagerMockRecorder) GenerateTurnToken() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateTurnToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateTurnToken))
+}
+
+// GetWGKey mocks base method.
+func (m *MockSecretsManager) GetWGKey() (wgtypes.Key, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetWGKey")
+	ret0, _ := ret[0].(wgtypes.Key)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetWGKey indicates an expected call of GetWGKey.
+func (mr *MockSecretsManagerMockRecorder) GetWGKey() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWGKey", reflect.TypeOf((*MockSecretsManager)(nil).GetWGKey))
+}
+
+// SetupRefresh mocks base method.
+func (m *MockSecretsManager) SetupRefresh(ctx context.Context, accountID, peerKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "SetupRefresh", ctx, accountID, peerKey)
+}
+
+// SetupRefresh indicates an expected call of SetupRefresh.
+func (mr *MockSecretsManagerMockRecorder) SetupRefresh(ctx, accountID, peerKey any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetupRefresh", reflect.TypeOf((*MockSecretsManager)(nil).SetupRefresh), ctx, accountID, peerKey)
+}
@@ -6,6 +6,14 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 )

+//go:generate go tool mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
+type Debouncer interface {
+	Stop()
+	TimerChannel() <-chan time.Time
+	ProcessUpdate(update *network_map.UpdateMessage) bool
+	GetPendingUpdates() []*network_map.UpdateMessage
+}
+
 // UpdateDebouncer implements a backpressure mechanism that:
 // - Sends the first update immediately
 // - Coalesces rapid subsequent network map updates (only latest matters)
@@ -0,0 +1,96 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./update_debouncer.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	reflect "reflect"
+	time "time"
+
+	network_map "github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockDebouncer is a mock of Debouncer interface.
+type MockDebouncer struct {
+	ctrl     *gomock.Controller
+	recorder *MockDebouncerMockRecorder
+	isgomock struct{}
+}
+
+// MockDebouncerMockRecorder is the mock recorder for MockDebouncer.
+type MockDebouncerMockRecorder struct {
+	mock *MockDebouncer
+}
+
+// NewMockDebouncer creates a new mock instance.
+func NewMockDebouncer(ctrl *gomock.Controller) *MockDebouncer {
+	mock := &MockDebouncer{ctrl: ctrl}
+	mock.recorder = &MockDebouncerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockDebouncer) EXPECT() *MockDebouncerMockRecorder {
+	return m.recorder
+}
+
+// GetPendingUpdates mocks base method.
+func (m *MockDebouncer) GetPendingUpdates() []*network_map.UpdateMessage {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPendingUpdates")
+	ret0, _ := ret[0].([]*network_map.UpdateMessage)
+	return ret0
+}
+
+// GetPendingUpdates indicates an expected call of GetPendingUpdates.
+func (mr *MockDebouncerMockRecorder) GetPendingUpdates() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPendingUpdates", reflect.TypeOf((*MockDebouncer)(nil).GetPendingUpdates))
+}
+
+// ProcessUpdate mocks base method.
+func (m *MockDebouncer) ProcessUpdate(update *network_map.UpdateMessage) bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ProcessUpdate", update)
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// ProcessUpdate indicates an expected call of ProcessUpdate.
+func (mr *MockDebouncerMockRecorder) ProcessUpdate(update any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProcessUpdate", reflect.TypeOf((*MockDebouncer)(nil).ProcessUpdate), update)
+}
+
+// Stop mocks base method.
+func (m *MockDebouncer) Stop() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Stop")
+}
+
+// Stop indicates an expected call of Stop.
+func (mr *MockDebouncerMockRecorder) Stop() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockDebouncer)(nil).Stop))
+}
+
+// TimerChannel mocks base method.
+func (m *MockDebouncer) TimerChannel() <-chan time.Time {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "TimerChannel")
+	ret0, _ := ret[0].(<-chan time.Time)
+	return ret0
+}
+
+// TimerChannel indicates an expected call of TimerChannel.
+func (mr *MockDebouncerMockRecorder) TimerChannel() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimerChannel", reflect.TypeOf((*MockDebouncer)(nil).TimerChannel))
+}