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

     1  package cohort2
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  	"sync"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/onflow/flow-go/model/flow"
    12  	"github.com/onflow/flow-go/model/libp2p/message"
    13  	mockcomponent "github.com/onflow/flow-go/module/component/mock"
    14  	"github.com/onflow/flow-go/network"
    15  	"github.com/onflow/flow-go/network/channels"
    16  	"github.com/onflow/flow-go/network/internal/testutils"
    17  )
    18  
    19  // EchoEngine is a simple engine that is used for testing the correctness of
    20  // driving the engines with libp2p, in addition to receiving and storing incoming messages
    21  // it also echos them back
    22  type EchoEngine struct {
    23  	sync.RWMutex
    24  	t        *testing.T
    25  	con      network.Conduit                  // used to directly communicate with the network
    26  	originID flow.Identifier                  // used to keep track of the id of the sender of the messages
    27  	event    chan interface{}                 // used to keep track of the events that the node receives
    28  	channel  chan channels.Channel            // used to keep track of the channels that events are received on
    29  	received chan struct{}                    // used as an indicator on reception of messages for testing
    30  	echomsg  string                           // used as a fix string to be included in the reply echos
    31  	seen     map[string]int                   // used to track the seen events
    32  	echo     bool                             // used to enable or disable echoing back the recvd message
    33  	send     testutils.ConduitSendWrapperFunc // used to provide play and plug wrapper around its conduit
    34  	mockcomponent.Component
    35  }
    36  
    37  func NewEchoEngine(t *testing.T, net network.EngineRegistry, cap int, channel channels.Channel, echo bool, send testutils.ConduitSendWrapperFunc) *EchoEngine {
    38  	te := &EchoEngine{
    39  		t:        t,
    40  		echomsg:  "this is an echo",
    41  		event:    make(chan interface{}, cap),
    42  		channel:  make(chan channels.Channel, cap),
    43  		received: make(chan struct{}, cap),
    44  		seen:     make(map[string]int),
    45  		echo:     echo,
    46  		send:     send,
    47  	}
    48  
    49  	c2, err := net.Register(channel, te)
    50  	require.NoError(te.t, err)
    51  	te.con = c2
    52  
    53  	return te
    54  }
    55  
    56  // SubmitLocal is implemented for a valid type assertion to Engine
    57  // any call to it fails the test
    58  func (te *EchoEngine) SubmitLocal(event interface{}) {
    59  	require.Fail(te.t, "not implemented")
    60  }
    61  
    62  // Submit is implemented for a valid type assertion to Engine
    63  // any call to it fails the test
    64  func (te *EchoEngine) Submit(channel channels.Channel, originID flow.Identifier, event interface{}) {
    65  	go func() {
    66  		err := te.Process(channel, originID, event)
    67  		if err != nil {
    68  			require.Fail(te.t, "could not process submitted event")
    69  		}
    70  	}()
    71  }
    72  
    73  // ProcessLocal is implemented for a valid type assertion to Engine
    74  // any call to it fails the test
    75  func (te *EchoEngine) ProcessLocal(event interface{}) error {
    76  	require.Fail(te.t, "not implemented")
    77  	return fmt.Errorf(" unexpected method called")
    78  }
    79  
    80  // Process receives an originID and an event and casts them into the corresponding fields of the
    81  // EchoEngine. It then flags the received channel on reception of an event.
    82  // It also sends back an echo of the message to the origin ID
    83  func (te *EchoEngine) Process(channel channels.Channel, originID flow.Identifier, event interface{}) error {
    84  	te.Lock()
    85  	defer te.Unlock()
    86  	te.originID = originID
    87  	te.event <- event
    88  	te.channel <- channel
    89  	te.received <- struct{}{}
    90  
    91  	// asserting event as string
    92  	lip2pEvent, ok := (event).(*message.TestMessage)
    93  	require.True(te.t, ok, "could not assert event as TestMessage")
    94  
    95  	// checks for duplication
    96  	strEvent := lip2pEvent.Text
    97  
    98  	// marks event as seen
    99  	te.seen[strEvent]++
   100  
   101  	// avoids endless circulation of echos by filtering echoing back the echo messages
   102  	if strings.HasPrefix(strEvent, te.echomsg) {
   103  		return nil
   104  	}
   105  
   106  	// do not echo back if not needed
   107  	if !te.echo {
   108  		return nil
   109  	}
   110  
   111  	// sends a echo back
   112  	msg := &message.TestMessage{
   113  		Text: fmt.Sprintf("%s: %s", te.echomsg, strEvent),
   114  	}
   115  
   116  	// Sends the message through the conduit's send wrapper
   117  	// Note: instead of directly interacting with the conduit, we use
   118  	// the wrapper function here, to enable this same implementation
   119  	// gets tested with different Conduit methods, i.e., publish, multicast, and unicast.
   120  	err := te.send(msg, te.con, originID)
   121  	// we allow one dial failure on echo due to connection tear down
   122  	// specially when the test is for a single message and not echo
   123  	// the sender side may close the connection as it does not expect any echo
   124  	if err != nil && !strings.Contains(err.Error(), "failed to dial") {
   125  		require.Fail(te.t, fmt.Sprintf("could not submit echo back to network: %s", err))
   126  	}
   127  	return nil
   128  }