github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/requester/discovery/identity_test.go (about)

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/filecoin-project/bacalhau/pkg/libp2p"
     8  	"github.com/filecoin-project/bacalhau/pkg/model"
     9  	"github.com/filecoin-project/bacalhau/pkg/transport/bprotocol"
    10  	"github.com/libp2p/go-libp2p/core/host"
    11  	"github.com/libp2p/go-libp2p/core/network"
    12  	"github.com/stretchr/testify/suite"
    13  )
    14  
    15  type IdentityNodeDiscovererSuite struct {
    16  	suite.Suite
    17  	discoverer *IdentityNodeDiscoverer
    18  	node1      host.Host
    19  	node2      host.Host
    20  	random3    host.Host
    21  }
    22  
    23  func (s *IdentityNodeDiscovererSuite) SetupSuite() {
    24  	// create 3 nodes network all connected to the first one as a bootstrap node,
    25  	// where only the first two nodes are implementing bprotcol
    26  	ctx := context.Background()
    27  	var err error
    28  	s.node1, err = libp2p.NewHostForTest(ctx)
    29  	s.NoError(err)
    30  	s.node2, err = libp2p.NewHostForTest(ctx)
    31  	s.NoError(err)
    32  	s.random3, err = libp2p.NewHostForTest(ctx)
    33  	s.NoError(err)
    34  
    35  	// implement bprotocol
    36  	s.node1.SetStreamHandler(bprotocol.AskForBidProtocolID, func(s network.Stream) {})
    37  	s.node2.SetStreamHandler(bprotocol.AskForBidProtocolID, func(s network.Stream) {})
    38  
    39  	// connect all nodes after implementing bprotocol to share the info in the initial handshake
    40  	s.NoError(s.node2.Connect(ctx, *host.InfoFromHost(s.node1)))
    41  	s.NoError(s.random3.Connect(ctx, *host.InfoFromHost(s.node1)))
    42  
    43  	s.discoverer = NewIdentityNodeDiscoverer(IdentityNodeDiscovererParams{
    44  		Host: s.node1,
    45  	})
    46  }
    47  
    48  func (s *IdentityNodeDiscovererSuite) TearDownSuite() {
    49  	s.NoError(s.node1.Close())
    50  	s.NoError(s.node2.Close())
    51  	s.NoError(s.random3.Close())
    52  }
    53  
    54  func TestIdentityNodeDiscovererSuite(t *testing.T) {
    55  	suite.Run(t, new(IdentityNodeDiscovererSuite))
    56  }
    57  
    58  func (s *IdentityNodeDiscovererSuite) TestFindNodes() {
    59  	discoverer := NewIdentityNodeDiscoverer(IdentityNodeDiscovererParams{
    60  		Host: s.node1,
    61  	})
    62  	peerIDs, err := discoverer.FindNodes(context.Background(), model.Job{})
    63  	s.NoError(err)
    64  
    65  	peerIDStrings := make([]string, len(peerIDs))
    66  	for i, p := range peerIDs {
    67  		peerIDStrings[i] = p.PeerInfo.ID.String()
    68  	}
    69  	s.ElementsMatch([]string{s.node1.ID().String(), s.node2.ID().String()}, peerIDStrings)
    70  }