From 8568fbffddb0fc09d1f73a336eaa00c656ef9862 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Wed, 24 Jul 2024 18:01:43 +0200
Subject: [PATCH] Fix test

---
 client/internal/peer/status.go       |  2 --
 client/internal/peer/worker_relay.go |  4 +---
 relay/client/client_test.go          | 28 ++++++++++++++++++++--------
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/client/internal/peer/status.go b/client/internal/peer/status.go
index f5f9a07a3..57accd28d 100644
--- a/client/internal/peer/status.go
+++ b/client/internal/peer/status.go
@@ -23,8 +23,6 @@ type State struct {
 	PubKey                     string
 	FQDN                       string
 	ConnStatus                 ConnStatus
-	connStatusRelay            ConnStatus
-	connStatusICE              ConnStatus
 	ConnStatusUpdate           time.Time
 	Relayed                    bool
 	LocalIceCandidateType      string
diff --git a/client/internal/peer/worker_relay.go b/client/internal/peer/worker_relay.go
index 2b90430f0..138f4aeae 100644
--- a/client/internal/peer/worker_relay.go
+++ b/client/internal/peer/worker_relay.go
@@ -8,7 +8,6 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/iface"
 	relayClient "github.com/netbirdio/netbird/relay/client"
 )
 
@@ -32,7 +31,6 @@ type WorkerRelay struct {
 	parentCtx    context.Context
 	log          *log.Entry
 	config       ConnConfig
-	wgInterface  iface.IWGIface
 	relayManager relayClient.ManagerService
 	conn         WorkerRelayCallbacks
 
@@ -120,7 +118,7 @@ func (w *WorkerRelay) wgStateCheck(conn net.Conn) {
 				w.conn.OnDisconnected()
 				return
 			}
-			resetTime := (lastHandshake.Add(wgHandshakeOvertime + wgHandshakePeriod)).Sub(time.Now())
+			resetTime := time.Until(lastHandshake.Add(wgHandshakeOvertime + wgHandshakePeriod))
 			timer.Reset(resetTime)
 		case <-w.ctx.Done():
 			return
diff --git a/relay/client/client_test.go b/relay/client/client_test.go
index b8685568b..7e511d302 100644
--- a/relay/client/client_test.go
+++ b/relay/client/client_test.go
@@ -200,10 +200,13 @@ func transfer(t *testing.T, testData []byte, peerPairs int) {
 
 	var transferDuration []time.Duration
 	wg := sync.WaitGroup{}
+	var writeErr error
+	var readErr error
 	for i := 0; i < len(connsSender); i++ {
 		wg.Add(2)
 		start := time.Now()
 		go func(i int) {
+			defer wg.Done()
 			pieceSize := 1024
 			testDataLen := len(testData)
 
@@ -212,34 +215,43 @@ func transfer(t *testing.T, testData []byte, peerPairs int) {
 				if end > testDataLen {
 					end = testDataLen
 				}
-				_, err := connsSender[i].Write(testData[j:end])
-				if err != nil {
-					t.Fatalf("failed to write to channel: %s", err)
+				_, writeErr = connsSender[i].Write(testData[j:end])
+				if writeErr != nil {
+					return
 				}
 			}
-			wg.Done()
+
 		}(i)
 
 		go func(i int, start time.Time) {
+			defer wg.Done()
 			buf := make([]byte, 8192)
 			rcv := 0
+			var n int
 			for receivedSize := 0; receivedSize < len(testData); {
 
-				n, err := connsReceiver[i].Read(buf)
-				if err != nil {
-					t.Fatalf("failed to read from channel: %s", err)
+				n, readErr = connsReceiver[i].Read(buf)
+				if readErr != nil {
+					return
 				}
 
 				receivedSize += n
 				rcv += n
 			}
 			transferDuration = append(transferDuration, time.Since(start))
-			wg.Done()
 		}(i, start)
 	}
 
 	wg.Wait()
 
+	if writeErr != nil {
+		t.Fatalf("failed to write to channel: %s", err)
+	}
+
+	if readErr != nil {
+		t.Fatalf("failed to read from channel: %s", err)
+	}
+
 	// calculate the megabytes per second from the average transferDuration against the dataSize
 	var totalDuration time.Duration
 	for _, d := range transferDuration {