github.com/koko1123/flow-go-1@v0.29.6/engine/consensus/dkg/messaging_engine_test.go (about)

     1  package dkg
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/rs/zerolog"
     8  	"github.com/stretchr/testify/mock"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	msg "github.com/koko1123/flow-go-1/model/messages"
    12  	"github.com/koko1123/flow-go-1/module/dkg"
    13  	module "github.com/koko1123/flow-go-1/module/mock"
    14  	"github.com/koko1123/flow-go-1/network/channels"
    15  	"github.com/koko1123/flow-go-1/network/mocknetwork"
    16  	"github.com/koko1123/flow-go-1/utils/unittest"
    17  )
    18  
    19  // Helper function to initialise an engine.
    20  func createTestEngine(t *testing.T) *MessagingEngine {
    21  	// setup mock conduit
    22  	conduit := &mocknetwork.Conduit{}
    23  	network := new(mocknetwork.Network)
    24  	network.On("Register", mock.Anything, mock.Anything).
    25  		Return(conduit, nil).
    26  		Once()
    27  
    28  	// setup local with nodeID
    29  	nodeID := unittest.IdentifierFixture()
    30  	me := new(module.Local)
    31  	me.On("NodeID").Return(nodeID)
    32  
    33  	engine, err := NewMessagingEngine(
    34  		zerolog.Logger{},
    35  		network,
    36  		me,
    37  		dkg.NewBrokerTunnel(),
    38  	)
    39  	require.NoError(t, err)
    40  
    41  	return engine
    42  }
    43  
    44  // TestForwardOutgoingMessages checks that the engine correctly forwards
    45  // outgoing messages from the tunnel's Out channel to the network conduit.
    46  func TestForwardOutgoingMessages(t *testing.T) {
    47  	// sender engine
    48  	engine := createTestEngine(t)
    49  
    50  	// expected DKGMessage
    51  	destinationID := unittest.IdentifierFixture()
    52  	expectedMsg := msg.NewDKGMessage(
    53  		[]byte("hello"),
    54  		"dkg-123",
    55  	)
    56  
    57  	// override the conduit to check that the Unicast call matches the expected
    58  	// message and destination ID
    59  	conduit := &mocknetwork.Conduit{}
    60  	conduit.On("Unicast", &expectedMsg, destinationID).
    61  		Return(nil).
    62  		Once()
    63  	engine.conduit = conduit
    64  
    65  	engine.tunnel.SendOut(msg.PrivDKGMessageOut{
    66  		DKGMessage: expectedMsg,
    67  		DestID:     destinationID,
    68  	})
    69  
    70  	time.Sleep(5 * time.Millisecond)
    71  
    72  	conduit.AssertExpectations(t)
    73  }
    74  
    75  // TestForwardIncomingMessages checks that the engine correclty forwards
    76  // messages from the conduit to the tunnel's In channel.
    77  func TestForwardIncomingMessages(t *testing.T) {
    78  	// sender engine
    79  	e := createTestEngine(t)
    80  
    81  	originID := unittest.IdentifierFixture()
    82  	expectedMsg := msg.PrivDKGMessageIn{
    83  		DKGMessage: msg.NewDKGMessage([]byte("hello"), "dkg-123"),
    84  		OriginID:   originID,
    85  	}
    86  
    87  	// launch a background routine to capture messages forwarded to the tunnel's
    88  	// In channel
    89  	doneCh := make(chan struct{})
    90  	go func() {
    91  		receivedMsg := <-e.tunnel.MsgChIn
    92  		require.Equal(t, expectedMsg, receivedMsg)
    93  		close(doneCh)
    94  	}()
    95  
    96  	err := e.Process(channels.DKGCommittee, originID, &expectedMsg.DKGMessage)
    97  	require.NoError(t, err)
    98  
    99  	unittest.RequireCloseBefore(t, doneCh, time.Second, "message not received")
   100  }