(fixup) Generate EnvelopeSender mock and use it for improving test

evidence/reactor_test.go

@@ -28,6 +28,11 @@ import (
 	"github.com/Finschia/ostracon/types"
 )
 
+type Peer struct {
+	*p2pmocks.Peer
+	*p2pmocks.EnvelopeSender
+}
+
 var (
 	numEvidence = 10
 	timeout     = 120 * time.Second // ridiculously high because CircleCI is slow
@@ -200,24 +205,24 @@ func TestReactorBroadcastEvidenceMemoryLeak(t *testing.T) {
 	pool, err := evidence.NewPool(evidenceDB, stateStore, blockStore)
 	require.NoError(t, err)
 
-	p := &p2pmocks.Peer{}
+	p := &Peer{Peer: &p2pmocks.Peer{}, EnvelopeSender: &p2pmocks.EnvelopeSender{}}
 
-	p.On("IsRunning").Once().Return(true)
-	p.On("IsRunning").Return(false)
+	p.Peer.On("IsRunning").Once().Return(true)
+	p.Peer.On("IsRunning").Return(false)
 	// check that we are not leaking any go-routines
 	// i.e. broadcastEvidenceRoutine finishes when peer is stopped
 	defer leaktest.CheckTimeout(t, 10*time.Second)()
 
-	p.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+	p.EnvelopeSender.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
 		e, ok := i.(p2p.Envelope)
 		return ok && e.ChannelID == evidence.EvidenceChannel
 	})).Return(false)
 	quitChan := make(<-chan struct{})
-	p.On("Quit").Return(quitChan)
+	p.Peer.On("Quit").Return(quitChan)
 	ps := peerState{2}
-	p.On("Get", types.PeerStateKey).Return(ps)
-	p.On("ID").Return("ABC")
-	p.On("String").Return("mock")
+	p.Peer.On("Get", types.PeerStateKey).Return(ps)
+	p.Peer.On("ID").Return("ABC")
+	p.Peer.On("String").Return("mock")
 
 	r := evidence.NewReactor(pool, config.P2P.RecvAsync, config.P2P.EvidenceRecvBufSize)
 	r.SetLogger(log.TestingLogger())

p2p/peer.go

@@ -16,6 +16,7 @@ import (
 )
 
 //go:generate ../scripts/mockery_generate.sh Peer
+//go:generate ../scripts/mockery_generate.sh EnvelopeSender
 
 const metricsTickerDuration = 10 * time.Second
 

statesync/reactor_test.go

@@ -26,6 +26,11 @@ import (
 	"github.com/Finschia/ostracon/version"
 )
 
+type Peer struct {
+	*p2pmocks.Peer
+	*p2pmocks.EnvelopeSender
+}
+
 func TestReactor_Receive_ChunkRequest(t *testing.T) {
 	testcases := map[string]struct {
 		request        *ssproto.ChunkRequest
@@ -59,11 +64,11 @@ func TestReactor_Receive_ChunkRequest(t *testing.T) {
 			}).Return(&abci.ResponseLoadSnapshotChunk{Chunk: tc.chunk}, nil)
 
 			// Mock peer to store response, if found
-			peer := &p2pmocks.Peer{}
-			peer.On("ID").Return(p2p.ID("id"))
+			peer := &Peer{Peer: &p2pmocks.Peer{}, EnvelopeSender: &p2pmocks.EnvelopeSender{}}
+			peer.Peer.On("ID").Return(p2p.ID("id"))
 			var response *ssproto.ChunkResponse
 			if tc.expectResponse != nil {
-				peer.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+				peer.EnvelopeSender.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
 					e, ok := i.(p2p.Envelope)
 					return ok && e.ChannelID == ChunkChannel
 				})).Run(func(args mock.Arguments) {
@@ -98,7 +103,8 @@ func TestReactor_Receive_ChunkRequest(t *testing.T) {
 			assert.Equal(t, tc.expectResponse, response)
 
 			conn.AssertExpectations(t)
-			peer.AssertExpectations(t)
+			peer.Peer.AssertExpectations(t)
+			peer.EnvelopeSender.AssertExpectations(t)
 		})
 	}
 }
@@ -150,10 +156,10 @@ func TestReactor_Receive_SnapshotsRequest(t *testing.T) {
 
 			// Mock peer to catch responses and store them in a slice
 			responses := []*ssproto.SnapshotsResponse{}
-			peer := &p2pmocks.Peer{}
+			peer := &Peer{Peer: &p2pmocks.Peer{}, EnvelopeSender: &p2pmocks.EnvelopeSender{}}
 			if len(tc.expectResponses) > 0 {
-				peer.On("ID").Return(p2p.ID("id"))
-				peer.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
+				peer.Peer.On("ID").Return(p2p.ID("id"))
+				peer.EnvelopeSender.On("SendEnvelope", mock.MatchedBy(func(i interface{}) bool {
 					e, ok := i.(p2p.Envelope)
 					return ok && e.ChannelID == SnapshotChannel
 				})).Run(func(args mock.Arguments) {
@@ -188,7 +194,8 @@ func TestReactor_Receive_SnapshotsRequest(t *testing.T) {
 			assert.Equal(t, tc.expectResponses, responses)
 
 			conn.AssertExpectations(t)
-			peer.AssertExpectations(t)
+			peer.Peer.AssertExpectations(t)
+			peer.EnvelopeSender.AssertExpectations(t)
 		})
 	}
 }