From e7283a81985e27b2e7f4045b0b1981270d47b732 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Wed, 28 Jan 2026 17:39:25 +0100
Subject: [PATCH] Update unit tests

---
 client/internal/peer/wg_watcher_test.go | 12 ++++++------
 shared/management/client/mock.go        |  9 +++++++++
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/client/internal/peer/wg_watcher_test.go b/client/internal/peer/wg_watcher_test.go
index e73d35668..3ce91cd46 100644
--- a/client/internal/peer/wg_watcher_test.go
+++ b/client/internal/peer/wg_watcher_test.go
@@ -35,11 +35,11 @@ func TestWGWatcher_EnableWgWatcher(t *testing.T) {
 	defer cancel()
 
 	onDisconnected := make(chan struct{}, 1)
-	go watcher.EnableWgWatcher(ctx,, func() {
+	go watcher.EnableWgWatcher(ctx, time.Now(), func() {
 		mlog.Infof("onDisconnectedFn")
 		onDisconnected <- struct{}{}
-	}, func(elapsed float64) {
-		mlog.Infof("onHandshakeSuccess: %.3fs", elapsed)
+	}, func(when time.Time) {
+		mlog.Infof("onHandshakeSuccess: %v", when)
 	})
 
 	// wait for initial reading
@@ -66,7 +66,7 @@ func TestWGWatcher_ReEnable(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		watcher.EnableWgWatcher(ctx, func() {}, func(elapsed float64) {})
+		watcher.EnableWgWatcher(ctx, time.Now(), func() {}, func(when time.Time) {})
 	}()
 	cancel()
 
@@ -77,9 +77,9 @@ func TestWGWatcher_ReEnable(t *testing.T) {
 	defer cancel()
 
 	onDisconnected := make(chan struct{}, 1)
-	go watcher.EnableWgWatcher(ctx,, func() {
+	go watcher.EnableWgWatcher(ctx, time.Now(), func() {
 		onDisconnected <- struct{}{}
-	}, func(elapsed float64) {})
+	}, func(when time.Time) {})
 
 	time.Sleep(2 * time.Second)
 	mocWgIface.disconnect()
diff --git a/shared/management/client/mock.go b/shared/management/client/mock.go
index ac96f7b36..a85f5e3be 100644
--- a/shared/management/client/mock.go
+++ b/shared/management/client/mock.go
@@ -18,6 +18,7 @@ type MockClient struct {
 	LoginFunc                      func(serverKey wgtypes.Key, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
 	GetDeviceAuthorizationFlowFunc func(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error)
 	GetPKCEAuthorizationFlowFunc   func(serverKey wgtypes.Key) (*proto.PKCEAuthorizationFlow, error)
+	GetServerURLFunc               func() string
 	SyncMetaFunc                   func(sysInfo *system.Info) error
 	LogoutFunc                     func() error
 	JobFunc                        func(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error
@@ -88,6 +89,14 @@ func (m *MockClient) GetNetworkMap(_ *system.Info) (*proto.NetworkMap, error) {
 	return nil, nil
 }
 
+// GetServerURL mock implementation of GetServerURL from mgm.Client interface
+func (m *MockClient) GetServerURL() string {
+	if m.GetServerURLFunc == nil {
+		return ""
+	}
+	return m.GetServerURLFunc()
+}
+
 func (m *MockClient) SyncMeta(sysInfo *system.Info) error {
 	if m.SyncMetaFunc == nil {
 		return nil