handshake ack tests

pkg/p2p/libp2p/internal/handshake/handshake_test.go

@@ -63,7 +63,7 @@ func TestHandshake(t *testing.T) {
 		stream1 := &StreamMock{readBuffer: &buffer1, writeBuffer: &buffer2}
 		stream2 := &StreamMock{readBuffer: &buffer2, writeBuffer: &buffer1}
 
-		w, _ := protobuf.NewWriterAndReader(stream2)
+		w, r := protobuf.NewWriterAndReader(stream2)
 		if err := w.WriteMsg(&pb.ShakeHand{
 			Address:   expectedInfo.Address,
 			NetworkID: expectedInfo.NetworkID,
@@ -80,6 +80,10 @@ func TestHandshake(t *testing.T) {
 		if *res != expectedInfo {
 			t.Fatalf("got %+v, expected %+v", res, info)
 		}
+
+		if err := r.ReadMsg(&pb.Ack{}); err != nil {
+			t.Fatal(err)
+		}
 	})
 
 	t.Run("ERROR - write error ", func(t *testing.T) {
@@ -142,6 +146,10 @@ func TestHandle(t *testing.T) {
 			t.Fatal(err)
 		}
 
+		if err := w.WriteMsg(&pb.Ack{}); err != nil {
+			t.Fatal(err)
+		}
+
 		res, err := handshakeService.Handle(stream1)
 		if err != nil {
 			t.Fatal(err)