Disconnect nodes in rollback_to

test/functional/feature_loan_payback_with_collateral.py

@@ -36,17 +36,6 @@ def set_test_params(self):
             '-simulatemainnet=1'
             ]]
 
-    def rollback_to(self, block):
-        self.log.info("rollback to: %d", block)
-        node = self.nodes[0]
-        current_height = node.getblockcount()
-        if current_height == block:
-            return
-        blockhash = node.getblockhash(block + 1)
-        node.invalidateblock(blockhash)
-        node.clearmempool()
-        assert_equal(block, node.getblockcount())
-
     def createOracles(self):
         self.oracle_address1 = self.nodes[0].getnewaddress("", "legacy")
         price_feeds = [{"currency": "USD", "token": "DFI"},

test/functional/feature_negative_interest.py

@@ -19,16 +19,6 @@ def getDecimalAmount(amount):
 
 
 class NegativeInterestTest (DefiTestFramework):
-    def rollback_to(self, block):
-        self.log.info("rollback to: %d", block)
-        current_height = self.nodes[0].getblockcount()
-        if current_height == block:
-            return
-        blockhash = self.nodes[0].getblockhash(block + 1)
-        self.nodes[0].invalidateblock(blockhash)
-        self.nodes[0].clearmempool()
-        assert_equal(block, self.nodes[0].getblockcount())
-
     def set_test_params(self):
         self.num_nodes = 1
         self.setup_clean_chain = True

test/functional/feature_on_chain_government_fee_distribution.py

@@ -7,9 +7,7 @@
 
 from test_framework.test_framework import DefiTestFramework
 from test_framework.util import (
-    assert_equal,
-    connect_nodes_bi,
-    disconnect_nodes,
+    assert_equal
 )
 from decimal import ROUND_DOWN, Decimal
 
@@ -120,19 +118,7 @@ def test_cfp_fee_distribution(self, amount, expectedFee, burnPct, vote, cycles=2
             history = self.nodes[0].listaccounthistory(mn3['ownerAuthAddress'], {"txtype": "ProposalFeeRedistribution"})
             assert_equal(history, [])
 
-        # Disconnect nodes and check connection count
-        for i in range(self.num_nodes - 1):
-            disconnect_nodes(self.nodes[i], i + 1)
-            assert_equal(self.nodes[i].getconnectioncount(), 0)
-        assert_equal(self.nodes[3].getconnectioncount(), 0)
-
-        # Rollback nodes in isolation
-        for i in range(self.num_nodes):
-            self.rollback_to(height, nodes=[i])
-
-        # Connect nodes
-        for i in range(self.num_nodes - 1):
-            connect_nodes_bi(self.nodes, i, i + 1)
+        self.rollback_to(height, nodes=[0, 1, 2, 3])
 
     def setup(self):
         # Get MN addresses

test/functional/rpc_getstoredinterest.py

@@ -43,16 +43,6 @@ def set_test_params(self):
             '-jellyfish_regtest=1', '-txindex=1', '-simulatemainnet=1']
         ]
 
-    # Utils
-    def rollback_to(self, block):
-        node = self.nodes[0]
-        current_height = node.getblockcount()
-        if current_height == block:
-            return
-        blockhash = node.getblockhash(block + 1)
-        node.invalidateblock(blockhash)
-        node.clearmempool()
-        assert_equal(block, node.getblockcount())
 
     def new_vault(self, loan_scheme, deposit=10):
         vaultId = self.nodes[0].createvault(self.account0, loan_scheme)

test/functional/test_framework/test_framework.py

@@ -15,6 +15,7 @@
 import sys
 import tempfile
 import time
+import re
 
 from typing import List
 from .authproxy import JSONRPCException
@@ -26,6 +27,7 @@
     PortSeed,
     assert_equal,
     check_json_precision,
+    connect_nodes,
     connect_nodes_bi,
     disconnect_nodes,
     get_datadir_path,
@@ -421,8 +423,25 @@ def rollback_to(self, block, nodes=None):
         if nodes is None:
             self._rollback_to(block)
         else:
+            connections = {}
             for node in nodes:
-                self._rollback_to(block, node=node)
+                nodes_connections = []
+                for x in self.nodes[node].getpeerinfo():
+                    if not x['inbound']:
+                        node_number = re.findall(r'\d+', x['subver'])[-1]
+                        nodes_connections.append(int(node_number))
+                connections[node] = nodes_connections
+
+            for node in nodes:
+                for x in connections[node]:
+                    disconnect_nodes(self.nodes[node], x)
+
+            for node in nodes:
+                self._rollback_to(block, node)
+
+            for node in nodes:
+                for x in connections[node]:
+                    connect_nodes(self.nodes[node], x)
 
     def run_test(self):
         """Tests must override this method to define test logic"""