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

     1  package network
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/stretchr/testify/mock"
     7  
     8  	"github.com/onflow/flow-go/model/flow"
     9  	"github.com/onflow/flow-go/network"
    10  	"github.com/onflow/flow-go/network/channels"
    11  	"github.com/onflow/flow-go/network/mocknetwork"
    12  )
    13  
    14  type EngineProcessFunc func(channels.Channel, flow.Identifier, interface{}) error
    15  type PublishFunc func(channels.Channel, interface{}, ...flow.Identifier) error
    16  
    17  // Conduit represents a mock conduit.
    18  
    19  // Network represents a mock network. The implementation is not concurrency-safe.
    20  type Network struct {
    21  	mocknetwork.Network
    22  	conduits    map[channels.Channel]*Conduit
    23  	engines     map[channels.Channel]network.MessageProcessor
    24  	publishFunc PublishFunc
    25  }
    26  
    27  var _ network.EngineRegistry = (*Network)(nil)
    28  
    29  // NewNetwork returns a new mock network.
    30  func NewNetwork() *Network {
    31  	return &Network{
    32  		Network:  mocknetwork.Network{},
    33  		conduits: make(map[channels.Channel]*Conduit),
    34  		engines:  make(map[channels.Channel]network.MessageProcessor),
    35  	}
    36  }
    37  
    38  // Register registers an engine with this mock network. If an engine is already registered on the
    39  // given channel, this will return an error.
    40  func (n *Network) Register(channel channels.Channel, engine network.MessageProcessor) (network.Conduit, error) {
    41  	_, ok := n.engines[channel]
    42  	if ok {
    43  		return nil, fmt.Errorf("channel already registered: %s", channel)
    44  	}
    45  
    46  	n.engines[channel] = engine
    47  	conduit := &Conduit{net: n, channel: channel}
    48  	n.conduits[channel] = conduit
    49  
    50  	return conduit, nil
    51  }
    52  
    53  // Send sends a message to the engine registered to the given channel on this mock network and returns
    54  // an error if one occurs. If no engine is registered, this is a noop.
    55  func (n *Network) Send(channel channels.Channel, originID flow.Identifier, event interface{}) error {
    56  	if eng, ok := n.engines[channel]; ok {
    57  		return eng.Process(channel, originID, event)
    58  	}
    59  	return nil
    60  }
    61  
    62  // OnPublish specifies the callback that should be executed when `Publish` is called on any conduits
    63  // created by this mock network.
    64  func (n *Network) OnPublish(publishFunc PublishFunc) *Network {
    65  	n.publishFunc = publishFunc
    66  	return n
    67  }
    68  
    69  // Engine represents a mock engine. The implementation is not concurrency-safe.
    70  type Engine struct {
    71  	mocknetwork.Engine
    72  }
    73  
    74  // NewEngine returns a new mock engine.
    75  func NewEngine() *Engine {
    76  	return &Engine{
    77  		mocknetwork.Engine{},
    78  	}
    79  }
    80  
    81  // OnProcess specifies the callback that should be executed when `Process` is called on this mock engine.
    82  func (e *Engine) OnProcess(processFunc EngineProcessFunc) *Engine {
    83  	e.On("Process", mock.AnythingOfType("channels.Channel"), mock.AnythingOfType("flow.Identifier"), mock.Anything).
    84  		Return((func(channels.Channel, flow.Identifier, interface{}) error)(processFunc))
    85  
    86  	return e
    87  }