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

     1  package ranking
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/filecoin-project/bacalhau/pkg/model"
     8  	"github.com/filecoin-project/bacalhau/pkg/requester"
     9  	"github.com/libp2p/go-libp2p/core/peer"
    10  	"github.com/stretchr/testify/suite"
    11  )
    12  
    13  type ChainSuite struct {
    14  	suite.Suite
    15  	chain   *Chain
    16  	peerID1 model.NodeInfo
    17  	peerID2 model.NodeInfo
    18  	peerID3 model.NodeInfo
    19  }
    20  
    21  func (s *ChainSuite) 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 *ChainSuite) SetupTest() {
    28  	s.chain = NewChain()
    29  }
    30  
    31  func TestChainSuite(t *testing.T) {
    32  	suite.Run(t, new(ChainSuite))
    33  }
    34  
    35  func (s *ChainSuite) TestRankNodes() {
    36  	s.chain.Add(newFixedRanker(0, 10, 10))
    37  	s.chain.Add(newFixedRanker(0, 0, 0))
    38  	s.chain.Add(newFixedRanker(0, 10, 20))
    39  
    40  	ranks, err := s.chain.RankNodes(context.Background(), model.Job{}, []model.NodeInfo{s.peerID1, s.peerID2, s.peerID3})
    41  	s.NoError(err)
    42  	s.Equal(3, len(ranks))
    43  	assertEquals(s.T(), ranks, "peerID1", 0)
    44  	assertEquals(s.T(), ranks, "peerID2", 20)
    45  	assertEquals(s.T(), ranks, "peerID3", 30)
    46  }
    47  
    48  func (s *ChainSuite) TestRankNodes_Negative() {
    49  	s.chain.Add(newFixedRanker(10, 10, 100))
    50  	s.chain.Add(newFixedRanker(0, 0, -1))
    51  	s.chain.Add(newFixedRanker(0, 10, 1000))
    52  
    53  	ranks, err := s.chain.RankNodes(context.Background(), model.Job{}, []model.NodeInfo{s.peerID1, s.peerID2, s.peerID3})
    54  	s.NoError(err)
    55  	s.Equal(3, len(ranks))
    56  	assertEquals(s.T(), ranks, "peerID1", 10)
    57  	assertEquals(s.T(), ranks, "peerID2", 20)
    58  	assertEquals(s.T(), ranks, "peerID3", -1)
    59  }
    60  
    61  func (s *ChainSuite) TestRankNodes_AllNegative() {
    62  	s.chain.Add(newFixedRanker(-99, -99, -99))
    63  	s.chain.Add(newFixedRanker(-1, -1, -1))
    64  	s.chain.Add(newFixedRanker(-999, 999, 999))
    65  
    66  	ranks, err := s.chain.RankNodes(context.Background(), model.Job{}, []model.NodeInfo{s.peerID1, s.peerID2, s.peerID3})
    67  	s.NoError(err)
    68  	s.Equal(3, len(ranks))
    69  	assertEquals(s.T(), ranks, "peerID1", -1)
    70  	assertEquals(s.T(), ranks, "peerID2", -1)
    71  	assertEquals(s.T(), ranks, "peerID3", -1)
    72  }
    73  
    74  // node Ranker that always returns the same set of nodes
    75  type fixedRanker struct {
    76  	ranks []int
    77  }
    78  
    79  func newFixedRanker(ranks ...int) *fixedRanker {
    80  	return &fixedRanker{
    81  		ranks: ranks,
    82  	}
    83  }
    84  
    85  func (f *fixedRanker) RankNodes(ctx context.Context, job model.Job, nodes []model.NodeInfo) ([]requester.NodeRank, error) {
    86  	ranks := make([]requester.NodeRank, len(nodes))
    87  	for i, rank := range f.ranks {
    88  		ranks[i] = requester.NodeRank{
    89  			NodeInfo: nodes[i],
    90  			Rank:     rank,
    91  		}
    92  	}
    93  	return ranks, nil
    94  }