From adbcd0c450ebc1631818637ebcb9435d0e48de79 Mon Sep 17 00:00:00 2001
From: William Banfield <wbanfield@gmail.com>
Date: Wed, 19 Oct 2022 14:50:35 -0400
Subject: [PATCH] fix statesync test

---
 p2p/switch_test.go        |  1 +
 statesync/reactor.go      |  1 -
 statesync/reactor_test.go | 15 ++++++++++++---
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/p2p/switch_test.go b/p2p/switch_test.go
index 83b983f5a..80ab3f663 100644
--- a/p2p/switch_test.go
+++ b/p2p/switch_test.go
@@ -469,6 +469,7 @@ func TestSwitchStopPeerForError(t *testing.T) {
 	p := sw1.Peers().List()[0]
 	e := Envelope{
 		ChannelID: 0x1,
+		Message:   &p2p.Message{},
 	}
 	p.Send(e)
 
diff --git a/statesync/reactor.go b/statesync/reactor.go
index 3eb8af31e..e30c234aa 100644
--- a/statesync/reactor.go
+++ b/statesync/reactor.go
@@ -130,7 +130,6 @@ func (r *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 			for _, snapshot := range snapshots {
 				r.Logger.Debug("Advertising snapshot", "height", snapshot.Height,
 					"format", snapshot.Format, "peer", src.ID())
-				m := mustEncodeMsg()
 				e := p2p.Envelope{
 					ChannelID: chID,
 					Message: toWrappedProto(&ssproto.SnapshotsResponse{
diff --git a/statesync/reactor_test.go b/statesync/reactor_test.go
index 053b47ef5..fa970d66a 100644
--- a/statesync/reactor_test.go
+++ b/statesync/reactor_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cosmos/gogoproto/proto"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -53,10 +54,18 @@ func TestReactor_Receive_ChunkRequest(t *testing.T) {
 			peer.On("ID").Return(p2p.ID("id"))
 			var response *ssproto.ChunkResponse
 			if tc.expectResponse != nil {
-				peer.On("Send", ChunkChannel, mock.Anything).Run(func(args mock.Arguments) {
-					msg, err := decodeMsg(args[1].([]byte))
+				peer.On("Send", mock.MatchedBy(func(i interface{}) bool {
+					e, ok := i.(p2p.Envelope)
+					return ok && e.ChannelID == ChunkChannel
+				})).Run(func(args mock.Arguments) {
+					e := args[0].(p2p.Envelope)
+
+					// Marshal to simulate a wire roundtrip.
+					bz, err := proto.Marshal(e.Message)
 					require.NoError(t, err)
-					response = msg.(*ssproto.ChunkResponse)
+					err = proto.Unmarshal(bz, e.Message)
+					require.NoError(t, err)
+					response = e.Message.(*ssproto.Message).GetChunkResponse()
 				}).Return(true)
 			}