github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/routing/inmemory/inmemory_test.go (about)

     1  package inmemory
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/filecoin-project/bacalhau/pkg/model"
     9  	"github.com/filecoin-project/bacalhau/pkg/requester"
    10  	"github.com/libp2p/go-libp2p/core/peer"
    11  	"github.com/stretchr/testify/suite"
    12  )
    13  
    14  type InMemoryNodeInfoStoreSuite struct {
    15  	suite.Suite
    16  	store *NodeInfoStore
    17  }
    18  
    19  func (s *InMemoryNodeInfoStoreSuite) SetupTest() {
    20  	s.store = NewNodeInfoStore(NodeInfoStoreParams{
    21  		TTL: 1 * time.Hour,
    22  	})
    23  }
    24  
    25  func TestInMemoryNodeInfoStoreSuite(t *testing.T) {
    26  	suite.Run(t, new(InMemoryNodeInfoStoreSuite))
    27  }
    28  
    29  func (s *InMemoryNodeInfoStoreSuite) Test_Get() {
    30  	ctx := context.Background()
    31  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
    32  	nodeInfo2 := generateNodeInfo("node2", model.EngineWasm)
    33  	s.NoError(s.store.Add(ctx, nodeInfo1))
    34  	s.NoError(s.store.Add(ctx, nodeInfo2))
    35  
    36  	// test Get
    37  	res1, err := s.store.Get(ctx, nodeInfo1.PeerInfo.ID)
    38  	s.NoError(err)
    39  	s.Equal(nodeInfo1, res1)
    40  
    41  	res2, err := s.store.Get(ctx, nodeInfo2.PeerInfo.ID)
    42  	s.NoError(err)
    43  	s.Equal(nodeInfo2, res2)
    44  }
    45  
    46  func (s *InMemoryNodeInfoStoreSuite) Test_GetNotFound() {
    47  	ctx := context.Background()
    48  	_, err := s.store.Get(ctx, peer.ID("node1"))
    49  	s.Error(err)
    50  	s.IsType(requester.ErrNodeNotFound{}, err)
    51  
    52  }
    53  
    54  func (s *InMemoryNodeInfoStoreSuite) Test_List() {
    55  	ctx := context.Background()
    56  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
    57  	nodeInfo2 := generateNodeInfo("node2", model.EngineWasm)
    58  	s.NoError(s.store.Add(ctx, nodeInfo1))
    59  	s.NoError(s.store.Add(ctx, nodeInfo2))
    60  
    61  	// test List
    62  	allNodeInfos, err := s.store.List(ctx)
    63  	s.NoError(err)
    64  	s.ElementsMatch([]model.NodeInfo{nodeInfo1, nodeInfo2}, allNodeInfos)
    65  }
    66  
    67  func (s *InMemoryNodeInfoStoreSuite) Test_ListForEngine() {
    68  	ctx := context.Background()
    69  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
    70  	nodeInfo2 := generateNodeInfo("node2", model.EngineWasm)
    71  	nodeInfo3 := generateNodeInfo("node3", model.EngineDocker, model.EngineWasm)
    72  	s.NoError(s.store.Add(ctx, nodeInfo1))
    73  	s.NoError(s.store.Add(ctx, nodeInfo2))
    74  	s.NoError(s.store.Add(ctx, nodeInfo3))
    75  
    76  	dockerNodes, err := s.store.ListForEngine(ctx, model.EngineDocker)
    77  	s.NoError(err)
    78  	s.ElementsMatch([]model.NodeInfo{nodeInfo1, nodeInfo3}, dockerNodes)
    79  
    80  	wasmNodes, err := s.store.ListForEngine(ctx, model.EngineWasm)
    81  	s.NoError(err)
    82  	s.ElementsMatch([]model.NodeInfo{nodeInfo2, nodeInfo3}, wasmNodes)
    83  }
    84  
    85  func (s *InMemoryNodeInfoStoreSuite) Test_Delete() {
    86  	ctx := context.Background()
    87  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
    88  	nodeInfo2 := generateNodeInfo("node2", model.EngineDocker, model.EngineWasm)
    89  	s.NoError(s.store.Add(ctx, nodeInfo1))
    90  	s.NoError(s.store.Add(ctx, nodeInfo2))
    91  
    92  	// delete first node
    93  	s.NoError(s.store.Delete(ctx, nodeInfo1.PeerInfo.ID))
    94  	dockerNodes, err := s.store.ListForEngine(ctx, model.EngineDocker)
    95  	s.NoError(err)
    96  	s.ElementsMatch([]model.NodeInfo{nodeInfo2}, dockerNodes)
    97  
    98  	wasmNodes, err := s.store.ListForEngine(ctx, model.EngineWasm)
    99  	s.NoError(err)
   100  	s.ElementsMatch([]model.NodeInfo{nodeInfo2}, wasmNodes)
   101  
   102  	// delete second node
   103  	s.NoError(s.store.Delete(ctx, nodeInfo2.PeerInfo.ID))
   104  	dockerNodes, err = s.store.ListForEngine(ctx, model.EngineDocker)
   105  	s.NoError(err)
   106  	s.Empty(dockerNodes)
   107  
   108  	wasmNodes, err = s.store.ListForEngine(ctx, model.EngineWasm)
   109  	s.NoError(err)
   110  	s.Empty(wasmNodes)
   111  }
   112  
   113  func (s *InMemoryNodeInfoStoreSuite) Test_Replace() {
   114  	ctx := context.Background()
   115  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
   116  	s.NoError(s.store.Add(ctx, nodeInfo1))
   117  
   118  	nodeInfo2 := generateNodeInfo("node1", model.EngineWasm)
   119  	nodeInfo2.PeerInfo.ID = nodeInfo1.PeerInfo.ID
   120  	s.NoError(s.store.Add(ctx, nodeInfo2))
   121  
   122  	res, err := s.store.Get(ctx, nodeInfo1.PeerInfo.ID)
   123  	s.NoError(err)
   124  	s.Equal(nodeInfo2, res)
   125  
   126  	// test List
   127  	allNodeInfos, err := s.store.List(ctx)
   128  	s.NoError(err)
   129  	s.ElementsMatch([]model.NodeInfo{nodeInfo2}, allNodeInfos)
   130  
   131  	// test ListForEngine
   132  	dockerNodes, err := s.store.ListForEngine(ctx, model.EngineDocker)
   133  	s.NoError(err)
   134  	s.Empty(dockerNodes)
   135  
   136  	wasmNodes, err := s.store.ListForEngine(ctx, model.EngineWasm)
   137  	s.NoError(err)
   138  	s.ElementsMatch([]model.NodeInfo{nodeInfo2}, wasmNodes)
   139  }
   140  
   141  func (s *InMemoryNodeInfoStoreSuite) Test_Eviction() {
   142  	ttl := 1 * time.Second
   143  	s.store = NewNodeInfoStore(NodeInfoStoreParams{
   144  		TTL: ttl,
   145  	})
   146  	ctx := context.Background()
   147  	nodeInfo1 := generateNodeInfo("node1", model.EngineDocker)
   148  	s.NoError(s.store.Add(ctx, nodeInfo1))
   149  
   150  	// test Get
   151  	res, err := s.store.Get(ctx, nodeInfo1.PeerInfo.ID)
   152  	s.NoError(err)
   153  	s.Equal(nodeInfo1, res)
   154  
   155  	// wait for eviction
   156  	time.Sleep(ttl + 100*time.Millisecond)
   157  	_, err = s.store.Get(ctx, nodeInfo1.PeerInfo.ID)
   158  	s.Error(err)
   159  	s.IsType(requester.ErrNodeNotFound{}, err)
   160  }
   161  
   162  func generateNodeInfo(id string, engines ...model.Engine) model.NodeInfo {
   163  	return model.NodeInfo{
   164  		PeerInfo: peer.AddrInfo{
   165  			ID: peer.ID(id),
   166  		},
   167  		NodeType: model.NodeTypeCompute,
   168  		ComputeNodeInfo: model.ComputeNodeInfo{
   169  			ExecutionEngines: engines,
   170  		},
   171  	}
   172  }