From 2202e2834018f3545dbb68a21be27e3773505c10 Mon Sep 17 00:00:00 2001
From: Bartosz Oleaczek <bartosz.oleaczek@nordsec.com>
Date: Tue, 7 Jan 2025 10:27:38 +0100
Subject: [PATCH] review fixes

---
 nc/nc.go      | 9 +++++----
 nc/nc_test.go | 8 +++++++-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/nc/nc.go b/nc/nc.go
index 18c75fcb..cd9c6b64 100644
--- a/nc/nc.go
+++ b/nc/nc.go
@@ -116,7 +116,7 @@ func (MqttClientBuilder) Build(opts *mqtt.ClientOptions) mqtt.Client {
 	return mqtt.NewClient(opts)
 }
 
-type TimeFunc func(int) time.Duration
+type CalculateRetryDelayForAttempt func(attempt int) time.Duration
 
 // Client is a client for Notification center
 type Client struct {
@@ -128,7 +128,7 @@ type Client struct {
 	subjectErr        events.Publisher[error]
 	subjectPeerUpdate events.Publisher[[]string]
 	credsFetcher      CredentialsGetter
-	timeFunc          TimeFunc
+	retryDelayFunc    CalculateRetryDelayForAttempt
 
 	startMu          sync.Mutex
 	started          bool
@@ -150,7 +150,7 @@ func NewClient(
 		subjectErr:        subjectErr,
 		subjectPeerUpdate: subjectPeerUpdate,
 		credsFetcher:      credsFetcher,
-		timeFunc:          network.ExponentialBackoff,
+		retryDelayFunc:    network.ExponentialBackoff,
 	}
 }
 
@@ -246,6 +246,7 @@ func (c *Client) tryConnect(
 		select {
 		case managementChan <- credentials.ExpirationDate:
 		case <-ctx.Done():
+			return client, connectionState
 		}
 
 		opts := c.createClientOptions(credentials, managementChan, ctx)
@@ -309,7 +310,7 @@ func (c *Client) connectWithBackoff(client mqtt.Client,
 				client.Disconnect(0)
 			}
 			return client
-		case <-time.After(c.timeFunc(tries)):
+		case <-time.After(c.retryDelayFunc(tries)):
 		}
 	}
 
diff --git a/nc/nc_test.go b/nc/nc_test.go
index 5201aae7..2476f38c 100644
--- a/nc/nc_test.go
+++ b/nc/nc_test.go
@@ -204,6 +204,7 @@ func TestConnectionCancellation(t *testing.T) {
 		connectionErr       error
 		fetchCredentialsErr error
 		tokenTimeout        time.Duration // how long client will wait for connection to be established
+		delayBeforeCancel   time.Duration
 	}{
 		{
 			name: "connection success",
@@ -224,6 +225,10 @@ func TestConnectionCancellation(t *testing.T) {
 			name:         "cancel while waiting for connection",
 			tokenTimeout: 10 * time.Second,
 		},
+		{
+			name:              "delay before cancel",
+			delayBeforeCancel: 10 * time.Millisecond,
+		},
 	}
 
 	for _, test := range tests {
@@ -248,7 +253,7 @@ func TestConnectionCancellation(t *testing.T) {
 			subjectErr:        &subs.Subject[error]{},
 			subjectPeerUpdate: &subs.Subject[[]string]{},
 			credsFetcher:      credsFetcher,
-			timeFunc:          func(i int) time.Duration { return test.tokenTimeout },
+			retryDelayFunc:    func(i int) time.Duration { return test.tokenTimeout },
 		}
 
 		t.Run(test.name, func(t *testing.T) {
@@ -259,6 +264,7 @@ func TestConnectionCancellation(t *testing.T) {
 				connectedChan <- true
 			}()
 
+			time.Sleep(test.delayBeforeCancel)
 			cancelFunc()
 
 			select {