From a07910abcd580ed07187794cf0e1faf040bb4212 Mon Sep 17 00:00:00 2001
From: Antoine Riard <ariard@student.42.fr>
Date: Sat, 17 Oct 2020 11:20:43 -0400
Subject: [PATCH] test: Makes wtxidrelay support a generic P2PInterface option

Its usage is extended beyond p2p_segwit.py in next commit.
---
 test/functional/p2p_segwit.py         | 10 +---------
 test/functional/p2p_tx_download.py    |  4 ++--
 test/functional/test_framework/p2p.py |  7 +++++--
 3 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/test/functional/p2p_segwit.py b/test/functional/p2p_segwit.py
index 29735b0fb3e..acbe3ee226c 100755
--- a/test/functional/p2p_segwit.py
+++ b/test/functional/p2p_segwit.py
@@ -37,7 +37,6 @@ from test_framework.messages import (
     msg_tx,
     msg_block,
     msg_no_witness_tx,
-    msg_verack,
     ser_uint256,
     ser_vector,
     sha256,
@@ -148,7 +147,7 @@ def test_witness_block(node, p2p, block, accepted, with_witness=True, reason=Non
 
 class TestP2PConn(P2PInterface):
     def __init__(self, wtxidrelay=False):
-        super().__init__()
+        super().__init__(wtxidrelay=wtxidrelay)
         self.getdataset = set()
         self.last_wtxidrelay = []
         self.lastgetdata = []
@@ -159,13 +158,6 @@ class TestP2PConn(P2PInterface):
     def on_inv(self, message):
         pass
 
-    def on_version(self, message):
-        if self.wtxidrelay:
-            super().on_version(message)
-        else:
-            self.send_message(msg_verack())
-            self.nServices = message.nServices
-
     def on_getdata(self, message):
         self.lastgetdata = message.inv
         for inv in message.inv:
diff --git a/test/functional/p2p_tx_download.py b/test/functional/p2p_tx_download.py
index 16d9302db87..a666d6c8d18 100755
--- a/test/functional/p2p_tx_download.py
+++ b/test/functional/p2p_tx_download.py
@@ -30,8 +30,8 @@ import time
 
 
 class TestP2PConn(P2PInterface):
-    def __init__(self):
-        super().__init__()
+    def __init__(self, wtxidrelay=True):
+        super().__init__(wtxidrelay=wtxidrelay)
         self.tx_getdata_count = 0
 
     def on_getdata(self, message):
diff --git a/test/functional/test_framework/p2p.py b/test/functional/test_framework/p2p.py
index 6846d31221b..0827e906d6e 100755
--- a/test/functional/test_framework/p2p.py
+++ b/test/functional/test_framework/p2p.py
@@ -289,7 +289,7 @@ class P2PInterface(P2PConnection):
 
     Individual testcases should subclass this and override the on_* methods
     if they want to alter message handling behaviour."""
-    def __init__(self, support_addrv2=False):
+    def __init__(self, support_addrv2=False, wtxidrelay=True):
         super().__init__()
 
         # Track number of messages of each type received.
@@ -309,6 +309,9 @@ class P2PInterface(P2PConnection):
 
         self.support_addrv2 = support_addrv2
 
+        # If the peer supports wtxid-relay
+        self.wtxidrelay = wtxidrelay
+
     def peer_connect(self, *args, services=NODE_NETWORK|NODE_WITNESS, send_version=True, **kwargs):
         create_conn = super().peer_connect(*args, **kwargs)
 
@@ -394,7 +397,7 @@ class P2PInterface(P2PConnection):
 
     def on_version(self, message):
         assert message.nVersion >= MIN_VERSION_SUPPORTED, "Version {} received. Test framework only supports versions greater than {}".format(message.nVersion, MIN_VERSION_SUPPORTED)
-        if message.nVersion >= 70016:
+        if message.nVersion >= 70016 and self.wtxidrelay:
             self.send_message(msg_wtxidrelay())
         self.send_message(msg_verack())
         if self.support_addrv2: