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

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"math"
     6  	"testing"
     7  
     8  	"github.com/filecoin-project/bacalhau/pkg/model"
     9  	"github.com/filecoin-project/bacalhau/pkg/routing/inmemory"
    10  	"github.com/libp2p/go-libp2p/core/peer"
    11  	"github.com/multiformats/go-multiaddr"
    12  	"github.com/stretchr/testify/suite"
    13  )
    14  
    15  type StoreNodeDiscovererSuite struct {
    16  	suite.Suite
    17  	discoverer *StoreNodeDiscoverer
    18  	store      *inmemory.NodeInfoStore
    19  }
    20  
    21  func (s *StoreNodeDiscovererSuite) SetupTest() {
    22  	s.store = inmemory.NewNodeInfoStore(inmemory.NodeInfoStoreParams{
    23  		TTL: math.MaxInt64,
    24  	})
    25  	s.discoverer = NewStoreNodeDiscoverer(StoreNodeDiscovererParams{
    26  		Store: s.store,
    27  	})
    28  }
    29  
    30  func TestStoreNodeDiscovererSuite(t *testing.T) {
    31  	suite.Run(t, new(StoreNodeDiscovererSuite))
    32  }
    33  
    34  func (s *StoreNodeDiscovererSuite) TestFindNodes() {
    35  	ctx := context.Background()
    36  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
    37  	nodeInfo2 := generateNodeInfo("node2", model.EngineDocker, model.EngineWasm)
    38  	s.NoError(s.store.Add(ctx, nodeInfo1))
    39  	s.NoError(s.store.Add(ctx, nodeInfo2))
    40  
    41  	// both nodes are returned when asked for docker nodes
    42  	job := model.Job{}
    43  	job.Spec.Engine = model.EngineDocker
    44  	peerIDs, err := s.discoverer.FindNodes(context.Background(), job)
    45  	s.NoError(err)
    46  	s.ElementsMatch([]model.NodeInfo{nodeInfo1, nodeInfo2}, peerIDs)
    47  
    48  	// only node2 is returned when asked for noop nodes
    49  	job.Spec.Engine = model.EngineNoop
    50  	peerIDs, err = s.discoverer.FindNodes(context.Background(), job)
    51  	s.NoError(err)
    52  	s.Empty(peerIDs)
    53  }
    54  
    55  func (s *StoreNodeDiscovererSuite) TestFindNodes_Empty() {
    56  	peerIDs, err := s.discoverer.FindNodes(context.Background(), model.Job{})
    57  	s.NoError(err)
    58  	s.Empty(peerIDs)
    59  }
    60  
    61  func generateNodeInfo(id string, engines ...model.Engine) model.NodeInfo {
    62  	return model.NodeInfo{
    63  		PeerInfo: peer.AddrInfo{
    64  			ID: peer.ID(id),
    65  			Addrs: []multiaddr.Multiaddr{
    66  				multiaddr.StringCast("/ip4/0.0.0.0/tcp/1234"),
    67  			},
    68  		},
    69  		NodeType: model.NodeTypeCompute,
    70  		ComputeNodeInfo: model.ComputeNodeInfo{
    71  			ExecutionEngines: engines,
    72  		},
    73  	}
    74  }