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

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"github.com/filecoin-project/bacalhau/pkg/model"
     9  	"github.com/libp2p/go-libp2p/core/peer"
    10  	"github.com/stretchr/testify/suite"
    11  )
    12  
    13  type ChainedSuite struct {
    14  	suite.Suite
    15  	chain   *Chain
    16  	peerID1 model.NodeInfo
    17  	peerID2 model.NodeInfo
    18  	peerID3 model.NodeInfo
    19  }
    20  
    21  func (s *ChainedSuite) SetupSuite() {
    22  	s.peerID1 = model.NodeInfo{PeerInfo: peer.AddrInfo{ID: peer.ID("peerID1")}}
    23  	s.peerID2 = model.NodeInfo{PeerInfo: peer.AddrInfo{ID: peer.ID("peerID2")}}
    24  	s.peerID3 = model.NodeInfo{PeerInfo: peer.AddrInfo{ID: peer.ID("peerID3")}}
    25  }
    26  
    27  func (s *ChainedSuite) SetupTest() {
    28  	s.chain = NewChain(false) // don't ignore errors
    29  }
    30  
    31  func TestChainedSuite(t *testing.T) {
    32  	suite.Run(t, new(ChainedSuite))
    33  }
    34  
    35  func (s *ChainedSuite) TestFindNodes() {
    36  	s.chain.Add(newFixedDiscoverer(s.peerID1))
    37  	s.chain.Add(newFixedDiscoverer(s.peerID2))
    38  	s.chain.Add(newFixedDiscoverer(s.peerID3))
    39  
    40  	peerIDs, err := s.chain.FindNodes(context.Background(), model.Job{})
    41  	s.NoError(err)
    42  	s.ElementsMatch([]model.NodeInfo{s.peerID1, s.peerID2, s.peerID3}, peerIDs)
    43  }
    44  
    45  func (s *ChainedSuite) TestFindNodes_Overlap() {
    46  	s.chain.Add(newFixedDiscoverer(s.peerID1, s.peerID2))
    47  	s.chain.Add(newFixedDiscoverer(s.peerID2, s.peerID3))
    48  
    49  	peerIDs, err := s.chain.FindNodes(context.Background(), model.Job{})
    50  	s.NoError(err)
    51  	s.ElementsMatch([]model.NodeInfo{s.peerID1, s.peerID2, s.peerID3}, peerIDs)
    52  }
    53  
    54  func (s *ChainedSuite) TestHandle_Error() {
    55  	s.chain.Add(newFixedDiscoverer(s.peerID1, s.peerID2))
    56  	s.chain.Add(newBadDiscoverer())
    57  	s.chain.Add(newFixedDiscoverer(s.peerID3))
    58  	peerIDs, err := s.chain.FindNodes(context.Background(), model.Job{})
    59  	s.Error(err)
    60  	s.Empty(peerIDs)
    61  }
    62  
    63  func (s *ChainedSuite) TestHandle_IgnoreError() {
    64  	s.chain.ignoreErrors = true
    65  	s.chain.Add(newFixedDiscoverer(s.peerID1, s.peerID2))
    66  	s.chain.Add(newBadDiscoverer())
    67  	s.chain.Add(newFixedDiscoverer(s.peerID3))
    68  
    69  	peerIDs, err := s.chain.FindNodes(context.Background(), model.Job{})
    70  	s.NoError(err)
    71  	s.ElementsMatch([]model.NodeInfo{s.peerID1, s.peerID2, s.peerID3}, peerIDs)
    72  }
    73  
    74  // node discoverer that always returns the same set of nodes
    75  type fixedDiscoverer struct {
    76  	peerIDs []model.NodeInfo
    77  }
    78  
    79  func newFixedDiscoverer(peerIDs ...model.NodeInfo) *fixedDiscoverer {
    80  	return &fixedDiscoverer{
    81  		peerIDs: peerIDs,
    82  	}
    83  }
    84  
    85  func (f *fixedDiscoverer) FindNodes(ctx context.Context, job model.Job) ([]model.NodeInfo, error) {
    86  	return f.peerIDs, nil
    87  }
    88  
    89  // node discoverer that always returns an error
    90  type badDiscoverer struct{}
    91  
    92  func newBadDiscoverer() *badDiscoverer {
    93  	return &badDiscoverer{}
    94  }
    95  
    96  func (b *badDiscoverer) FindNodes(ctx context.Context, job model.Job) ([]model.NodeInfo, error) {
    97  	return nil, errors.New("bad discoverer")
    98  }