github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/stub/hub.go (about)

     1  package stub
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/onflow/flow-go/model/flow"
    11  )
    12  
    13  // Hub is a test helper that mocks a network overlay.
    14  // It maintains a set of network instances and enables them to directly exchange message
    15  // over the memory.
    16  type Hub struct {
    17  	sync.RWMutex
    18  	networks map[flow.Identifier]*Network
    19  	Buffer   *Buffer
    20  }
    21  
    22  // NewNetworkHub creates and returns a new Hub instance.
    23  func NewNetworkHub() *Hub {
    24  	return &Hub{
    25  		networks: make(map[flow.Identifier]*Network),
    26  		Buffer:   NewBuffer(),
    27  	}
    28  }
    29  
    30  // DeliverAll delivers all the buffered messages in the Network instances attached to the Hub
    31  // to their destination.
    32  // Note that the delivery of messages is done in asynchronous mode, i.e., sender and receiver are
    33  // synchronized over delivery and not execution of the message.
    34  func (h *Hub) DeliverAll() {
    35  	for _, network := range h.networks {
    36  		network.DeliverAll(false)
    37  	}
    38  }
    39  
    40  // DeliverAllEventually attempts on delivery of all the buffered messages in the Network instances
    41  // attached to this instance of Hub. Once the delivery is done, it evaluates and returns the
    42  // condition function. It fails if delivery of all buffered messages in the Network instances
    43  // attached to this Hub is not getting done within 10 seconds.
    44  // Note that the delivery of messages is done in asynchronous mode, i.e., sender and receiver are
    45  // synchronized over delivery and not execution of the message.
    46  func (h *Hub) DeliverAllEventually(t *testing.T, condition func() bool) {
    47  	h.DeliverAllEventuallyUntil(t, condition, time.Second*10, time.Millisecond*500)
    48  }
    49  
    50  // DeliverAllEventuallyUntil attempts attempts on delivery of all the buffered messages in the Network instances
    51  // attached to this instance of Hub. Once the delivery is done, it evaluates and returns the
    52  // condition function. It fails if delivery of all buffered messages in the Network instances
    53  // attached to this Hub is not getting done within `waitFor` time interval.
    54  // It checks the status of message deliveries at every `tick` time interval.
    55  // Note that the delivery of messages is done in asynchronous mode, i.e., sender and receiver are
    56  // synchronized over delivery and not execution of the message.
    57  func (h *Hub) DeliverAllEventuallyUntil(t *testing.T, condition func() bool, waitFor time.Duration, tick time.Duration) {
    58  	require.Eventually(t, func() bool {
    59  		h.DeliverAll()
    60  		return condition()
    61  	}, waitFor, tick)
    62  }
    63  
    64  // GetNetwork returns the Network instance attached to the node ID.
    65  func (h *Hub) GetNetwork(nodeID flow.Identifier) (*Network, bool) {
    66  	h.RLock()
    67  	defer h.RUnlock()
    68  
    69  	net, ok := h.networks[nodeID]
    70  	return net, ok
    71  }
    72  
    73  // AddNetwork stores the reference of the Network in the Hub, in order for networks to find
    74  // other networks to send events directly.
    75  func (h *Hub) AddNetwork(net *Network) {
    76  	h.Lock()
    77  	defer h.Unlock()
    78  
    79  	h.networks[net.GetID()] = net
    80  }