From cb3a99b9ccdba9e26a0999aebce4f00c064024ee Mon Sep 17 00:00:00 2001
From: Gabriel mermelstein <gabrielmermelstein@gmail.com>
Date: Wed, 23 Oct 2024 14:58:07 +0300
Subject: [PATCH] cleanup

---
 wakuv2/nwaku.go      | 12 ++++++++----
 wakuv2/nwaku_test.go | 40 ++++++++++++++++++++--------------------
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/wakuv2/nwaku.go b/wakuv2/nwaku.go
index e06687c2a..bd5b8c738 100644
--- a/wakuv2/nwaku.go
+++ b/wakuv2/nwaku.go
@@ -2160,18 +2160,22 @@ func (self *Waku) wakuStoreQuery(
 	return "", errors.New(errMsg)
 }
 
-func (self *Waku) WakuPeerExchangeRequest(numPeers uint64) (string, error) {
+func (self *Waku) WakuPeerExchangeRequest(numPeers uint64) (uint64, error) {
 	var resp = C.allocResp()
 	defer C.freeResp(resp)
 
 	C.cGoWakuPeerExchangeQuery(self.wakuCtx, C.uint64_t(numPeers), resp)
 	if C.getRet(resp) == C.RET_OK {
-		msg := C.GoStringN(C.getMyCharPtr(resp), C.int(C.getMyCharLen(resp)))
-		return msg, nil
+		numRecvPeersStr := C.GoStringN(C.getMyCharPtr(resp), C.int(C.getMyCharLen(resp)))
+		numRecvPeers, err := strconv.ParseUint(numRecvPeersStr, 10, 64)
+		if err != nil {
+			return 0, err
+		}
+		return numRecvPeers, nil
 	}
 	errMsg := "error WakuPeerExchangeRequest: " +
 		C.GoStringN(C.getMyCharPtr(resp), C.int(C.getMyCharLen(resp)))
-	return "", errors.New(errMsg)
+	return 0, errors.New(errMsg)
 }
 
 func (self *Waku) WakuConnect(peerMultiAddr string, timeoutMs int) error {
diff --git a/wakuv2/nwaku_test.go b/wakuv2/nwaku_test.go
index 8db82a034..4ebab128b 100644
--- a/wakuv2/nwaku_test.go
+++ b/wakuv2/nwaku_test.go
@@ -6,7 +6,6 @@ package wakuv2
 import (
 	"context"
 	"errors"
-	"fmt"
 	"slices"
 	"testing"
 	"time"
@@ -192,9 +191,7 @@ func TestBasicWakuV2(t *testing.T) {
 
 	// Sanity check, not great, but it's probably helpful
 	err = tt.RetryWithBackOff(func() error {
-		
 		numConnected, err := w.GetNumConnectedPeers()
-		fmt.Println("numConnected: ", numConnected)
 		if err != nil {
 			return err
 		}
@@ -335,15 +332,13 @@ func TestPeerExchange(t *testing.T) {
 		Shards: []uint16{64},
 		PeerExchange: false,
 		Discv5UdpPort: 9001,
-		TcpPort: 60001,
+		TcpPort: 60010,
 	}
 
 	discV5Node, err := New(nil, "", &discV5NodeConfig, logger.Named("discV5Node"), nil, nil, nil, nil)
 	require.NoError(t, err)
 	require.NoError(t, discV5Node.Start())
 
-	fmt.Println("--------- GABRIEL started discV5Node")
-
 	time.Sleep(1 * time.Second)
 
 	discV5NodePeerId, err := discV5Node.PeerID()
@@ -352,7 +347,6 @@ func TestPeerExchange(t *testing.T) {
 	discv5NodeEnr, err := discV5Node.ENR()
 	require.NoError(t, err)
 	
-	
 	// start node which serves as PeerExchange server
 	pxServerConfig := WakuConfig{
 		EnableRelay: true,
@@ -363,25 +357,19 @@ func TestPeerExchange(t *testing.T) {
 		PeerExchange: true,
 		Discv5UdpPort: 9000,
 		Discv5BootstrapNodes: []string{discv5NodeEnr.String()},
-		TcpPort: 60002,
+		TcpPort: 60011,
 	}
 
 	pxServerNode, err := New(nil, "", &pxServerConfig, logger.Named("pxServerNode"), nil, nil, nil, nil)
 	require.NoError(t, err)
 	require.NoError(t, pxServerNode.Start())
 
-	fmt.Println("------ GABRIEL started big sleep")
-	time.Sleep(10 * time.Second)
-	fmt.Println("------ GABRIEL finished big sleep")
+	time.Sleep(1 * time.Second)
 
 	serverNodeMa, err := pxServerNode.ListenAddresses()
 	require.NoError(t, err)
 	require.NotNil(t, serverNodeMa)
 
-	discv5Peers, err := discV5Node.GetPeerIdsFromPeerStore()
-	require.NoError(t, err)
-	fmt.Println("------ GABRIEL discV5Node peers: ", discv5Peers)
-
 	// Sanity check, not great, but it's probably helpful
 	options := func(b *backoff.ExponentialBackOff) {
 		b.MaxElapsedTime = 30 * time.Second
@@ -390,9 +378,7 @@ func TestPeerExchange(t *testing.T) {
 	// Check that pxServerNode has discV5Node in its Peer Store
 	err = tt.RetryWithBackOff(func() error {
 		peers, err := pxServerNode.GetPeerIdsFromPeerStore()
-		
-		fmt.Println("------ GABRIEL pxServerNode peers: ", peers)
-		
+				
 		if err != nil {
 			return err
 		}
@@ -414,7 +400,7 @@ func TestPeerExchange(t *testing.T) {
 		Shards: []uint16{64},
 		PeerExchange: true,
 		Discv5UdpPort: 9002,
-		TcpPort: 60003,
+		TcpPort: 60012,
 		PeerExchangeNode: serverNodeMa[0].String(),
 	}
 
@@ -429,7 +415,6 @@ func TestPeerExchange(t *testing.T) {
 	
 	err = tt.RetryWithBackOff(func() error {
 		peers, err := lightNode.GetPeerIdsFromPeerStore()
-		fmt.Println("------ GABRIEL lightnode peers: ", peers)
 		if err != nil {
 			return err
 		}
@@ -441,6 +426,21 @@ func TestPeerExchange(t *testing.T) {
 	}, options)
 	require.NoError(t, err)
 
+	// Now perform the PX request manually to see if it also works
+	err = tt.RetryWithBackOff(func() error {
+		numPeersReceived, err := lightNode.WakuPeerExchangeRequest(1)
+		if err != nil {
+			return err
+		}
+
+		if numPeersReceived == 1 {
+			return nil
+		}
+		return errors.New("Peer Exchange is not returning peers")
+	}, options)
+	require.NoError(t, err)
+
+
 	require.NoError(t, lightNode.Stop())
 	require.NoError(t, pxServerNode.Stop())
 	require.NoError(t, discV5Node.Stop())