From a94e350ac0e5b65ef23a84b05fb10d1204c98c97 Mon Sep 17 00:00:00 2001
From: stratospher <44024636+stratospher@users.noreply.github.com>
Date: Sat, 5 Feb 2022 23:28:02 +0530
Subject: [PATCH] [test] Build v2 P2P messages

---
 test/functional/test_framework/p2p.py | 29 ++++++++++++++++++---------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/test/functional/test_framework/p2p.py b/test/functional/test_framework/p2p.py
index dcb93e9acce..4ababe01dc1 100755
--- a/test/functional/test_framework/p2p.py
+++ b/test/functional/test_framework/p2p.py
@@ -82,6 +82,7 @@ from test_framework.util import (
 )
 from test_framework.v2_p2p import (
     EncryptedP2PState,
+    MSGTYPE_TO_SHORTID,
     SHORTID,
 )
 
@@ -386,15 +387,25 @@ class P2PConnection(asyncio.Protocol):
         """Build a serialized P2P message"""
         msgtype = message.msgtype
         data = message.serialize()
-        tmsg = self.magic_bytes
-        tmsg += msgtype
-        tmsg += b"\x00" * (12 - len(msgtype))
-        tmsg += struct.pack("<I", len(data))
-        th = sha256(data)
-        h = sha256(th)
-        tmsg += h[:4]
-        tmsg += data
-        return tmsg
+        if self.supports_v2_p2p:
+            if msgtype in SHORTID.values():
+                tmsg = MSGTYPE_TO_SHORTID.get(msgtype).to_bytes(1, 'big')
+            else:
+                tmsg = b"\x00"
+                tmsg += msgtype
+                tmsg += b"\x00" * (12 - len(msgtype))
+            tmsg += data
+            return self.v2_state.v2_enc_packet(tmsg)
+        else:
+            tmsg = self.magic_bytes
+            tmsg += msgtype
+            tmsg += b"\x00" * (12 - len(msgtype))
+            tmsg += struct.pack("<I", len(data))
+            th = sha256(data)
+            h = sha256(th)
+            tmsg += h[:4]
+            tmsg += data
+            return tmsg
 
     def _log_message(self, direction, msg):
         """Logs a message being sent or received over the connection."""