diff --git a/itests/clients/node_client.go b/itests/clients/node_client.go index 028f97d00..1e6002694 100644 --- a/itests/clients/node_client.go +++ b/itests/clients/node_client.go @@ -284,8 +284,8 @@ func (c *NodesClients) SynchronizedWavesBalances( } func (c *NodesClients) Handshake() { - c.GoClient.Handshake() - c.ScalaClient.Handshake() + c.GoClient.Connection.SendHandshake() + c.ScalaClient.Connection.SendHandshake() } func (c *NodesClients) SendToNodes(t *testing.T, m proto.Message, scala bool) { @@ -300,8 +300,10 @@ func (c *NodesClients) SendToNodes(t *testing.T, m proto.Message, scala bool) { } func (c *NodesClients) Close(t *testing.T) { - c.GoClient.Close(t) - c.ScalaClient.Close(t) + c.GoClient.GRPCClient.Close(t) + c.GoClient.Connection.Close() + c.ScalaClient.GRPCClient.Close(t) + c.ScalaClient.Connection.Close() } func (c *NodesClients) requestNodesAvailableBalances( diff --git a/pkg/networking/session.go b/pkg/networking/session.go index 6cfc79e73..bdab4bdd4 100644 --- a/pkg/networking/session.go +++ b/pkg/networking/session.go @@ -348,7 +348,7 @@ func (s *Session) readMessagePayload(hdr Header, conn io.Reader) error { s.logger.Debug("Invoking OnReceive handler", "message", base64.StdEncoding.EncodeToString(s.receiveBuffer.Bytes())) } - s.config.handler.OnReceive(s, bytes.NewReader(s.receiveBuffer.Bytes())) // Invoke OnReceive handler. + s.config.handler.OnReceive(s, s.receiveBuffer) // Invoke OnReceive handler. return nil } diff --git a/pkg/networking/session_test.go b/pkg/networking/session_test.go index 745ee194b..c30bfac81 100644 --- a/pkg/networking/session_test.go +++ b/pkg/networking/session_test.go @@ -56,7 +56,7 @@ func TestSuccessfulSession(t *testing.T) { require.NoError(t, wErr) assert.Equal(t, 5, n) }) - sc2 := serverHandler.On("OnReceive", ss, bytes.NewReader(encodeMessage("Hello session"))). + sc2 := serverHandler.On("OnReceive", ss, bytes.NewBuffer(encodeMessage("Hello session"))). Once().Return() sc2.NotBefore(sc1). Run(func(_ mock.Arguments) { @@ -75,7 +75,7 @@ func TestSuccessfulSession(t *testing.T) { require.NoError(t, wErr) assert.Equal(t, 17, n) }) - cl2 := clientHandler.On("OnReceive", cs, bytes.NewReader(encodeMessage("Hi"))).Once().Return() + cl2 := clientHandler.On("OnReceive", cs, bytes.NewBuffer(encodeMessage("Hi"))).Once().Return() cl2.NotBefore(cl1). Run(func(_ mock.Arguments) { cWG.Done()