chore: improve tests to neighbor messages

lib/grandpa/mocks_generate_test.go

@@ -3,4 +3,4 @@
 
 package grandpa
 
-//go:generate mockgen -destination=mocks_test.go -package $GOPACKAGE . BlockState,GrandpaState
+//go:generate mockgen -destination=mocks_test.go -package $GOPACKAGE . BlockState,GrandpaState,Network

lib/grandpa/network_test.go

@@ -4,13 +4,17 @@
 package grandpa
 
 import (
+	"context"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/golang/mock/gomock"
 
 	"github.com/libp2p/go-libp2p-core/peer"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -78,57 +82,82 @@ func TestHandleNetworkMessage(t *testing.T) {
 	require.False(t, propagate)
 }
 
-func TestSendNeighbourMessage(t *testing.T) {
-	gs, st := newTestService(t)
-	go gs.notifyNeighbor(time.Second)
-
-	digest := types.NewDigest()
-	prd, err := types.NewBabeSecondaryPlainPreDigest(0, 1).ToPreRuntimeDigest()
-	require.NoError(t, err)
-	err = digest.Add(*prd)
-	require.NoError(t, err)
-	block := &types.Block{
-		Header: types.Header{
-			ParentHash: st.Block.GenesisHash(),
-			Number:     1,
-			Digest:     digest,
+func TestNotifyNeighbor(t *testing.T) {
+	const interval = 2 * time.Second
+
+	tests := map[string]struct {
+		notifyInterval     time.Duration
+		finalizeBlock      bool
+		finalizeBlockAfter time.Duration
+		expectWithin       time.Duration
+	}{
+		"should_send_neighbor_message": {
+			expectWithin:   2 * time.Second,
+			notifyInterval: interval,
+		},
+		"should_reset_timer_and_then_send_neighbor_message": {
+			finalizeBlock:      true,
+			finalizeBlockAfter: 1 * time.Second,
+			notifyInterval:     interval,
+			expectWithin:       3 * time.Second,
 		},
-		Body: types.Body{},
-	}
-
-	err = st.Block.AddBlock(block)
-	require.NoError(t, err)
-
-	hash := block.Header.Hash()
-	round := uint64(7)
-	setID := uint64(33)
-	err = st.Block.SetFinalisedHash(hash, round, setID)
-	require.NoError(t, err)
-
-	expected := &NeighbourMessage{
-		Version: 1,
-		SetID:   setID,
-		Round:   round,
-		Number:  1,
-	}
-
-	select {
-	case <-time.After(time.Second):
-		t.Fatal("did not send message")
-	case msg := <-gs.network.(*testNetwork).out:
-		nm, ok := msg.(*NeighbourMessage)
-		require.True(t, ok)
-		require.Equal(t, expected, nm)
 	}
 
-	require.Equal(t, expected, gs.neighbourMessage)
-
-	select {
-	case <-time.After(time.Second * 2):
-		t.Fatal("did not send message")
-	case msg := <-gs.network.(*testNetwork).out:
-		nm, ok := msg.(*NeighbourMessage)
-		require.True(t, ok)
-		require.Equal(t, expected, nm)
+	for tname, tt := range tests {
+		tt := tt
+		t.Run(tname, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			mockedNet := NewMockNetwork(ctrl)
+			mockedCh := make(chan *types.FinalisationInfo)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			s := Service{
+				ctx: ctx,
+				state: &State{
+					round: 1,
+					setID: 0,
+				},
+				finalisedCh: mockedCh,
+				head:        &types.Header{Number: 0},
+				network:     mockedNet,
+			}
+
+			expectedNeighborMessage := &NeighbourMessage{
+				Version: 1,
+				Round:   s.state.round,
+				SetID:   s.state.setID,
+				Number:  uint32(s.head.Number),
+			}
+			cm, err := expectedNeighborMessage.ToConsensusMessage()
+			require.NoError(t, err)
+
+			timecheck := new(time.Time)
+
+			wg := new(sync.WaitGroup)
+			wg.Add(1)
+
+			ensureGossipMessageCalledRightTime := func(_ network.NotificationsMessage) {
+				defer wg.Done()
+				const roundOverSec = 1 * time.Second
+
+				calledWithin := time.Now().Sub(*timecheck)
+				calledWithin = calledWithin.Round(roundOverSec) // avoid decimal points
+				assert.Equal(t, tt.expectWithin, calledWithin)
+			}
+
+			mockedNet.EXPECT().GossipMessage(cm).Times(1).DoAndReturn(ensureGossipMessageCalledRightTime)
+
+			*timecheck = time.Now()
+			go s.notifyNeighbor(tt.notifyInterval)
+
+			if tt.finalizeBlock {
+				<-time.After(tt.finalizeBlockAfter)
+				mockedCh <- &types.FinalisationInfo{}
+			}
+
+			wg.Wait()
+		})
 	}
 }