go-waku/waku/v2/peermanager/peer_manager_test.go

package peermanager

import (
	"context"
	"crypto/rand"
	"fmt"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p/core/host"
	libp2pProtocol "github.com/libp2p/go-libp2p/core/protocol"
	"github.com/multiformats/go-multiaddr"
	"github.com/stretchr/testify/require"
	"github.com/waku-org/go-waku/tests"
	wps "github.com/waku-org/go-waku/waku/v2/peerstore"
	"github.com/waku-org/go-waku/waku/v2/utils"
)

func getAddr(h host.Host) multiaddr.Multiaddr {
	id, _ := multiaddr.NewMultiaddr(fmt.Sprintf("/p2p/%s", h.ID().Pretty()))
	return h.Network().ListenAddresses()[0].Encapsulate(id)
}

func initTest(t *testing.T) (context.Context, *PeerManager, func()) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	// hosts
	h1, err := tests.MakeHost(ctx, 0, rand.Reader)
	require.NoError(t, err)
	defer h1.Close()

	// host 1 is used by peer manager
	pm := NewPeerManager(10, utils.Logger())
	pm.SetHost(h1)

	return ctx, pm, func() {
		cancel()
		h1.Close()
	}
}

func TestServiceSlots(t *testing.T) {
	ctx, pm, deferFn := initTest(t)
	defer deferFn()

	h2, err := tests.MakeHost(ctx, 0, rand.Reader)
	require.NoError(t, err)
	defer h2.Close()

	h3, err := tests.MakeHost(ctx, 0, rand.Reader)
	require.NoError(t, err)
	defer h3.Close()
	// protocols
	protocol := libp2pProtocol.ID("test/protocol")
	protocol1 := libp2pProtocol.ID("test/protocol1")

	// add h2 peer to peer manager
	t.Log(h2.ID())
	_, err = pm.AddPeer(getAddr(h2), wps.Static, libp2pProtocol.ID(protocol))
	require.NoError(t, err)

	///////////////
	// getting peer for protocol
	///////////////

	// select peer from pm, currently only h2 is set in pm
	peerID, err := pm.SelectPeer(protocol, nil)
	require.NoError(t, err)
	require.Equal(t, peerID, h2.ID())

	// add h3 peer to peer manager
	_, err = pm.AddPeer(getAddr(h3), wps.Static, libp2pProtocol.ID(protocol))
	require.NoError(t, err)

	// check that returned peer is h2 or h3 peer
	peerID, err = pm.SelectPeer(protocol, nil)
	require.NoError(t, err)
	if peerID == h2.ID() || peerID == h3.ID() {
		//Test success
		t.Log("Random peer selection per protocol successful")
	} else {
		t.FailNow()
	}

	///////////////
	// getting peer for protocol1
	///////////////
	h4, err := tests.MakeHost(ctx, 0, rand.Reader)
	require.NoError(t, err)
	defer h4.Close()

	_, err = pm.SelectPeer(protocol1, nil)
	require.Error(t, err, utils.ErrNoPeersAvailable)

	// add h4 peer for protocol1
	_, err = pm.AddPeer(getAddr(h4), wps.Static, libp2pProtocol.ID(protocol1))
	require.NoError(t, err)

	//Test peer selection for protocol1
	peerID, err = pm.SelectPeer(protocol1, nil)
	require.NoError(t, err)
	require.Equal(t, peerID, h4.ID())

}

func TestDefaultProtocol(t *testing.T) {
	ctx, pm, deferFn := initTest(t)
	defer deferFn()
	///////////////
	// check peer for default protocol
	///////////////
	//Test empty peer selection for relay protocol
	_, err := pm.SelectPeer(WakuRelayIDv200, nil)
	require.Error(t, err, utils.ErrNoPeersAvailable)

	///////////////
	// getting peer for default protocol
	///////////////
	h5, err := tests.MakeHost(ctx, 0, rand.Reader)
	require.NoError(t, err)
	defer h5.Close()

	//Test peer selection for relay protocol from peer store
	_, err = pm.AddPeer(getAddr(h5), wps.Static, WakuRelayIDv200)
	require.NoError(t, err)

	// since we are not passing peerList, selectPeer fn using filterByProto checks in PeerStore for peers with same protocol.
	peerID, err := pm.SelectPeer(WakuRelayIDv200, nil)
	require.NoError(t, err)
	require.Equal(t, peerID, h5.ID())
}

func TestAdditionAndRemovalOfPeer(t *testing.T) {
	ctx, pm, deferFn := initTest(t)
	defer deferFn()
	///////////////
	// set h6 peer for protocol2 and remove that peer and check again
	///////////////
	//Test random peer selection
	protocol2 := libp2pProtocol.ID("test/protocol2")
	h6, err := tests.MakeHost(ctx, 0, rand.Reader)
	require.NoError(t, err)
	defer h6.Close()

	_, err = pm.AddPeer(getAddr(h6), wps.Static, protocol2)
	require.NoError(t, err)

	peerID, err := pm.SelectPeer(protocol2, nil)
	require.NoError(t, err)
	require.Equal(t, peerID, h6.ID())

	pm.RemovePeer(peerID)
	_, err = pm.SelectPeer(protocol2, nil)
	require.Error(t, err, utils.ErrNoPeersAvailable)
}

func TestConnectToRelayPeers(t *testing.T) {

	ctx, pm, deferFn := initTest(t)
	pc, err := NewPeerConnectionStrategy(pm, 120*time.Second, pm.logger)
	require.NoError(t, err)
	pm.SetPeerConnector(pc)
	err = pc.Start(ctx)
	require.NoError(t, err)
	pm.Start(ctx)

	defer deferFn()

	pm.connectToRelayPeers()

}