Disconnect nodes in rollback_to

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/test_framework/test_framework.py

@@ -22,9 +22,11 @@
 from .mininode import NetworkThread
 from .util import (
     MAX_NODES,
+    PORT_MIN,
     PortSeed,
     assert_equal,
     check_json_precision,
+    connect_nodes,
     connect_nodes_bi,
     disconnect_nodes,
     get_datadir_path,
@@ -420,8 +422,27 @@ 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']:
+                        nodes_connections.append(int(x['addr'].split(':')[1]) - PORT_MIN)
+                connections.append(nodes)
+
+            for node in nodes:
+                for x in connections[node]:
+                    disconnect_nodes(self.nodes[node], x)
+
+            for node in nodes:
+                assert (self.nodes[node].getconnectioncount() == 0)
+
+            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"""