From bc242813ad0feedd4db846a44287bb850004c7d9 Mon Sep 17 00:00:00 2001
From: Gabriel mermelstein <gabrielmermelstein@gmail.com>
Date: Tue, 17 Dec 2024 18:28:34 +0100
Subject: [PATCH] chore: pass ctx in dial functions

---
 waku/nwaku.go      |  8 ++------
 waku/nwaku_test.go | 18 ++++++++++++------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/waku/nwaku.go b/waku/nwaku.go
index b6ab7d7..939f6da 100644
--- a/waku/nwaku.go
+++ b/waku/nwaku.go
@@ -417,10 +417,8 @@ func (w *Waku) ListenAddresses() ([]multiaddr.Multiaddr, error) {
 	return w.node.ListenAddresses()
 }
 
-func (w *Waku) DialPeer(address multiaddr.Multiaddr) error {
+func (w *Waku) DialPeer(ctx context.Context, address multiaddr.Multiaddr) error {
 	// Using WakuConnect so it matches the go-waku's behavior and terminology
-	ctx, cancel := context.WithTimeout(w.ctx, requestTimeout)
-	defer cancel()
 	return w.node.Connect(ctx, address)
 }
 
@@ -429,9 +427,7 @@ func (w *Waku) RelayPublish(ctx context.Context, message *pb.WakuMessage, pubsub
 	return w.node.RelayPublish(ctx, message, pubsubTopic)
 }
 
-func (w *Waku) DialPeerByID(peerID peer.ID, protocol libp2pproto.ID) error {
-	ctx, cancel := context.WithTimeout(w.ctx, requestTimeout)
-	defer cancel()
+func (w *Waku) DialPeerByID(ctx context.Context, peerID peer.ID, protocol libp2pproto.ID) error {
 	return w.node.DialPeerByID(ctx, peerID, protocol)
 }
 
diff --git a/waku/nwaku_test.go b/waku/nwaku_test.go
index 1301fdb..9d5ce18 100644
--- a/waku/nwaku_test.go
+++ b/waku/nwaku_test.go
@@ -103,7 +103,9 @@ func TestBasicWaku(t *testing.T) {
 	require.True(t, isDisconnected, "nwaku should be disconnected from the store node")
 
 	// Re-connect
-	err = w.DialPeer(storeNodeMa)
+	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
+	defer cancel()
+	err = w.DialPeer(ctx, storeNodeMa)
 	require.NoError(t, err)
 
 	// Check that we are connected again
@@ -384,7 +386,9 @@ func TestDial(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, receiverPeerCount == 0, "Receiver node should have no connected peers")
 	// Dial
-	err = dialerNode.DialPeer(receiverMultiaddr[0])
+	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
+	defer cancel()
+	err = dialerNode.DialPeer(ctx, receiverMultiaddr[0])
 	require.NoError(t, err)
 	time.Sleep(1 * time.Second)
 	// Check that both nodes now have one connected peer
@@ -438,7 +442,9 @@ func TestRelay(t *testing.T) {
 	require.NotNil(t, receiverMultiaddr)
 
 	// Dial so they become peers
-	err = senderNode.DialPeer(receiverMultiaddr[0])
+	ctx1, cancel1 := context.WithTimeout(context.Background(), requestTimeout)
+	defer cancel1()
+	err = senderNode.DialPeer(ctx1, receiverMultiaddr[0])
 	require.NoError(t, err)
 	time.Sleep(1 * time.Second)
 	// Check that both nodes now have one connected peer
@@ -457,9 +463,9 @@ func TestRelay(t *testing.T) {
 	}
 	// send message
 	pubsubTopic := FormatWakuRelayTopic(senderNodeWakuConfig.ClusterID, senderNodeWakuConfig.Shards[0])
-	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
-	defer cancel()
-	senderNode.RelayPublish(ctx, message, pubsubTopic)
+	ctx2, cancel2 := context.WithTimeout(context.Background(), requestTimeout)
+	defer cancel2()
+	senderNode.RelayPublish(ctx2, message, pubsubTopic)
 
 	wg := sync.WaitGroup{}
 	wg.Add(1)