github.com/onflow/flow-go@v0.33.17/engine/consensus/dkg/messaging_engine_test.go (about)

     1  package dkg
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/mock"
     9  	"github.com/stretchr/testify/require"
    10  	"github.com/stretchr/testify/suite"
    11  
    12  	msg "github.com/onflow/flow-go/model/messages"
    13  	"github.com/onflow/flow-go/module/dkg"
    14  	"github.com/onflow/flow-go/module/irrecoverable"
    15  	"github.com/onflow/flow-go/module/metrics"
    16  	mockmodule "github.com/onflow/flow-go/module/mock"
    17  	"github.com/onflow/flow-go/network/channels"
    18  	"github.com/onflow/flow-go/network/mocknetwork"
    19  	"github.com/onflow/flow-go/utils/unittest"
    20  )
    21  
    22  // MessagingEngineSuite encapsulates unit tests for the MessagingEngine.
    23  type MessagingEngineSuite struct {
    24  	suite.Suite
    25  
    26  	conduit *mocknetwork.Conduit
    27  	network *mocknetwork.Network
    28  	me      *mockmodule.Local
    29  
    30  	engine *MessagingEngine
    31  }
    32  
    33  func TestMessagingEngine(t *testing.T) {
    34  	suite.Run(t, new(MessagingEngineSuite))
    35  }
    36  
    37  func (ms *MessagingEngineSuite) SetupTest() {
    38  	// setup mock conduit
    39  	ms.conduit = mocknetwork.NewConduit(ms.T())
    40  	ms.network = mocknetwork.NewNetwork(ms.T())
    41  	ms.network.On("Register", mock.Anything, mock.Anything).
    42  		Return(ms.conduit, nil).
    43  		Once()
    44  
    45  	// setup local with nodeID
    46  	nodeID := unittest.IdentifierFixture()
    47  	ms.me = mockmodule.NewLocal(ms.T())
    48  	ms.me.On("NodeID").Return(nodeID).Maybe()
    49  
    50  	engine, err := NewMessagingEngine(
    51  		unittest.Logger(),
    52  		ms.network,
    53  		ms.me,
    54  		dkg.NewBrokerTunnel(),
    55  		metrics.NewNoopCollector(),
    56  		DefaultMessagingEngineConfig(),
    57  	)
    58  	require.NoError(ms.T(), err)
    59  	ms.engine = engine
    60  }
    61  
    62  // TestForwardOutgoingMessages checks that the engine correctly forwards
    63  // outgoing messages from the tunnel's Out channel to the network conduit.
    64  func (ms *MessagingEngineSuite) TestForwardOutgoingMessages() {
    65  	ctx, cancel := irrecoverable.NewMockSignalerContextWithCancel(ms.T(), context.Background())
    66  	ms.engine.Start(ctx)
    67  	defer cancel()
    68  
    69  	// expected DKGMessage
    70  	destinationID := unittest.IdentifierFixture()
    71  	expectedMsg := msg.NewDKGMessage(
    72  		[]byte("hello"),
    73  		"dkg-123",
    74  	)
    75  
    76  	done := make(chan struct{})
    77  	ms.conduit.On("Unicast", &expectedMsg, destinationID).
    78  		Run(func(_ mock.Arguments) { close(done) }).
    79  		Return(nil).
    80  		Once()
    81  
    82  	ms.engine.tunnel.SendOut(msg.PrivDKGMessageOut{
    83  		DKGMessage: expectedMsg,
    84  		DestID:     destinationID,
    85  	})
    86  
    87  	unittest.RequireCloseBefore(ms.T(), done, time.Second, "message not sent")
    88  }
    89  
    90  // TestForwardIncomingMessages checks that the engine correctly forwards
    91  // messages from the conduit to the tunnel's MsgChIn channel.
    92  func (ms *MessagingEngineSuite) TestForwardIncomingMessages() {
    93  	ctx, cancel := irrecoverable.NewMockSignalerContextWithCancel(ms.T(), context.Background())
    94  	ms.engine.Start(ctx)
    95  	defer cancel()
    96  
    97  	originID := unittest.IdentifierFixture()
    98  	expectedMsg := msg.PrivDKGMessageIn{
    99  		DKGMessage: msg.NewDKGMessage([]byte("hello"), "dkg-123"),
   100  		OriginID:   originID,
   101  	}
   102  
   103  	// launch a background routine to capture messages forwarded to the tunnel's MsgChIn channel
   104  	done := make(chan struct{})
   105  	go func() {
   106  		receivedMsg := <-ms.engine.tunnel.MsgChIn
   107  		require.Equal(ms.T(), expectedMsg, receivedMsg)
   108  		close(done)
   109  	}()
   110  
   111  	err := ms.engine.Process(channels.DKGCommittee, originID, &expectedMsg.DKGMessage)
   112  	require.NoError(ms.T(), err)
   113  
   114  	unittest.RequireCloseBefore(ms.T(), done, time.Second, "message not received")
   115  }