github.com/koko1123/flow-go-1@v0.29.6/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/koko1123/flow-go-1/model/flow"
     9  	"github.com/koko1123/flow-go-1/network"
    10  	"github.com/koko1123/flow-go-1/network/channels"
    11  	"github.com/koko1123/flow-go-1/network/mocknetwork"
    12  )
    13  
    14  type EngineProcessFunc func(channels.Channel, flow.Identifier, interface{}) error
    15  type NetworkPublishFunc func(channels.Channel, interface{}, ...flow.Identifier) error
    16  
    17  // Conduit represents a mock conduit.
    18  type Conduit struct {
    19  	mocknetwork.Conduit
    20  	net     *Network
    21  	channel channels.Channel
    22  }
    23  
    24  // Publish sends a message on this mock network, invoking any callback that has
    25  // been specified. This will panic if no callback is found.
    26  func (c *Conduit) Publish(event interface{}, targetIDs ...flow.Identifier) error {
    27  	if c.net.publishFunc != nil {
    28  		return c.net.publishFunc(c.channel, event, targetIDs...)
    29  	}
    30  	panic("Publish called but no callback function was found.")
    31  }
    32  
    33  // Network represents a mock network. The implementation is not concurrency-safe.
    34  type Network struct {
    35  	mocknetwork.Network
    36  	conduits    map[channels.Channel]*Conduit
    37  	engines     map[channels.Channel]network.MessageProcessor
    38  	publishFunc NetworkPublishFunc
    39  }
    40  
    41  // NewNetwork returns a new mock network.
    42  func NewNetwork() *Network {
    43  	return &Network{
    44  		Network:  mocknetwork.Network{},
    45  		conduits: make(map[channels.Channel]*Conduit),
    46  		engines:  make(map[channels.Channel]network.MessageProcessor),
    47  	}
    48  }
    49  
    50  // Register registers an engine with this mock network. If an engine is already registered on the
    51  // given channel, this will return an error.
    52  func (n *Network) Register(channel channels.Channel, engine network.MessageProcessor) (network.Conduit, error) {
    53  	_, ok := n.engines[channel]
    54  	if ok {
    55  		return nil, fmt.Errorf("channel already registered: %s", channel)
    56  	}
    57  
    58  	n.engines[channel] = engine
    59  	conduit := &Conduit{net: n, channel: channel}
    60  	n.conduits[channel] = conduit
    61  
    62  	return conduit, nil
    63  }
    64  
    65  // Send sends a message to the engine registered to the given channel on this mock network and returns
    66  // an error if one occurs. If no engine is registered, this is a noop.
    67  func (n *Network) Send(channel channels.Channel, originID flow.Identifier, event interface{}) error {
    68  	if eng, ok := n.engines[channel]; ok {
    69  		return eng.Process(channel, originID, event)
    70  	}
    71  	return nil
    72  }
    73  
    74  // OnPublish specifies the callback that should be executed when `Publish` is called on any conduits
    75  // created by this mock network.
    76  func (n *Network) OnPublish(publishFunc NetworkPublishFunc) *Network {
    77  	n.publishFunc = publishFunc
    78  	return n
    79  }
    80  
    81  // Engine represents a mock engine. The implementation is not concurrency-safe.
    82  type Engine struct {
    83  	mocknetwork.Engine
    84  }
    85  
    86  // NewEngine returns a new mock engine.
    87  func NewEngine() *Engine {
    88  	return &Engine{
    89  		mocknetwork.Engine{},
    90  	}
    91  }
    92  
    93  // OnProcess specifies the callback that should be executed when `Process` is called on this mock engine.
    94  func (e *Engine) OnProcess(processFunc EngineProcessFunc) *Engine {
    95  	e.On("Process", mock.AnythingOfType("channels.Channel"), mock.AnythingOfType("flow.Identifier"), mock.Anything).
    96  		Return((func(channels.Channel, flow.Identifier, interface{}) error)(processFunc))
    97  
    98  	return e
    99  }