fix: discv5 tests

Author: richard-ramos

waku/v2/discv5/discover_test.go

@@ -4,15 +4,20 @@ import (
 	"context"
 	"crypto/ecdsa"
 	"fmt"
+	"math"
+	"net"
+	"strconv"
 	"testing"
 	"time"
 
 	gcrypto "github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/multiformats/go-multiaddr"
 	"github.com/status-im/go-waku/tests"
 	"github.com/status-im/go-waku/waku/v2/utils"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 
 	"github.com/libp2p/go-libp2p"
 	libp2pcrypto "github.com/libp2p/go-libp2p-core/crypto"
@@ -41,25 +46,87 @@ func createHost(t *testing.T) (host.Host, int, *ecdsa.PrivateKey) {
 	return host, port, privKey
 }
 
+func newLocalnode(priv *ecdsa.PrivateKey, ipAddr *net.TCPAddr, udpPort int, wakuFlags utils.WakuEnrBitfield, advertiseAddr *net.IP, log *zap.Logger) (*enode.LocalNode, error) {
+	db, err := enode.OpenDB("")
+	if err != nil {
+		return nil, err
+	}
+	localnode := enode.NewLocalNode(db, priv)
+	localnode.SetFallbackUDP(udpPort)
+	localnode.Set(enr.WithEntry(utils.WakuENRField, wakuFlags))
+	localnode.SetFallbackIP(net.IP{127, 0, 0, 1})
+	localnode.SetStaticIP(ipAddr.IP)
+
+	if udpPort > 0 && udpPort <= math.MaxUint16 {
+		localnode.Set(enr.UDP(uint16(udpPort))) // lgtm [go/incorrect-integer-conversion]
+	} else {
+		log.Error("setting udpPort", zap.Int("port", udpPort))
+	}
+
+	if ipAddr.Port > 0 && ipAddr.Port <= math.MaxUint16 {
+		localnode.Set(enr.TCP(uint16(ipAddr.Port))) // lgtm [go/incorrect-integer-conversion]
+	} else {
+		log.Error("setting tcpPort", zap.Int("port", ipAddr.Port))
+	}
+
+	if advertiseAddr != nil {
+		localnode.SetStaticIP(*advertiseAddr)
+	}
+
+	return localnode, nil
+}
+
+func extractIP(addr multiaddr.Multiaddr) (*net.TCPAddr, error) {
+	ipStr, err := addr.ValueForProtocol(multiaddr.P_IP4)
+	if err != nil {
+		return nil, err
+	}
+
+	portStr, err := addr.ValueForProtocol(multiaddr.P_TCP)
+	if err != nil {
+		return nil, err
+	}
+	port, err := strconv.Atoi(portStr)
+	if err != nil {
+		return nil, err
+	}
+	return &net.TCPAddr{
+		IP:   net.ParseIP(ipStr),
+		Port: port,
+	}, nil
+}
+
 func TestDiscV5(t *testing.T) {
 	// Host1 <-> Host2 <-> Host3
 
+	// H1
 	host1, _, prvKey1 := createHost(t)
 	udpPort1, err := tests.FindFreePort(t, "127.0.0.1", 3)
 	require.NoError(t, err)
-	d1, err := NewDiscoveryV5(host1, host1.Addrs(), prvKey1, utils.NewWakuEnrBitfield(true, true, true, true), utils.Logger(), WithUDPPort(udpPort1))
+	ip1, _ := extractIP(host1.Addrs()[0])
+	l1, err := newLocalnode(prvKey1, ip1, udpPort1, utils.NewWakuEnrBitfield(true, true, true, true), nil, utils.Logger())
+	require.NoError(t, err)
+	d1, err := NewDiscoveryV5(host1, prvKey1, l1, utils.Logger(), WithUDPPort(udpPort1))
 	require.NoError(t, err)
 
+	// H2
 	host2, _, prvKey2 := createHost(t)
+	ip2, _ := extractIP(host2.Addrs()[0])
 	udpPort2, err := tests.FindFreePort(t, "127.0.0.1", 3)
 	require.NoError(t, err)
-	d2, err := NewDiscoveryV5(host2, host2.Addrs(), prvKey2, utils.NewWakuEnrBitfield(true, true, true, true), utils.Logger(), WithUDPPort(udpPort2), WithBootnodes([]*enode.Node{d1.localnode.Node()}))
+	l2, err := newLocalnode(prvKey2, ip2, udpPort2, utils.NewWakuEnrBitfield(true, true, true, true), nil, utils.Logger())
+	require.NoError(t, err)
+	d2, err := NewDiscoveryV5(host2, prvKey2, l2, utils.Logger(), WithUDPPort(udpPort2), WithBootnodes([]*enode.Node{d1.localnode.Node()}))
 	require.NoError(t, err)
 
+	// H3
 	host3, _, prvKey3 := createHost(t)
+	ip3, _ := extractIP(host3.Addrs()[0])
 	udpPort3, err := tests.FindFreePort(t, "127.0.0.1", 3)
 	require.NoError(t, err)
-	d3, err := NewDiscoveryV5(host3, host3.Addrs(), prvKey3, utils.NewWakuEnrBitfield(true, true, true, true), utils.Logger(), WithUDPPort(udpPort3), WithBootnodes([]*enode.Node{d2.localnode.Node()}))
+	l3, err := newLocalnode(prvKey3, ip3, udpPort3, utils.NewWakuEnrBitfield(true, true, true, true), nil, utils.Logger())
+	require.NoError(t, err)
+	d3, err := NewDiscoveryV5(host3, prvKey3, l3, utils.Logger(), WithUDPPort(udpPort3), WithBootnodes([]*enode.Node{d2.localnode.Node()}))
 	require.NoError(t, err)
 
 	defer d1.Stop()
@@ -75,7 +142,7 @@ func TestDiscV5(t *testing.T) {
 	err = d3.Start()
 	require.NoError(t, err)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
 	peerChan, err := d3.FindPeers(ctx, "", discovery.Limit(2))