From 80ff9c5b6398b4603ced64811267a4cc53576453 Mon Sep 17 00:00:00 2001
From: Michael Schmoock <michael@schmoock.net>
Date: Sun, 8 Mar 2020 11:36:25 +0100
Subject: [PATCH] test: refactor lockup_drain add helpers drain and
 force_feerates

---
 contrib/pyln-testing/pyln/testing/utils.py | 23 +++++++++++++++++
 tests/test_pay.py                          | 30 ++++++----------------
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/contrib/pyln-testing/pyln/testing/utils.py b/contrib/pyln-testing/pyln/testing/utils.py
index 828f92027..952fb71a3 100644
--- a/contrib/pyln-testing/pyln/testing/utils.py
+++ b/contrib/pyln-testing/pyln/testing/utils.py
@@ -1,4 +1,5 @@
 from bitcoin.rpc import RawProxy as BitcoinProxy
+from pyln.client import RpcError
 from pyln.testing.btcproxy import BitcoinRpcProxy
 from collections import OrderedDict
 from decimal import Decimal
@@ -811,6 +812,7 @@ class LightningNode(object):
                     if 'htlcs' in channel:
                         wait_for(lambda: len(self.rpc.listpeers()['peers'][p]['channels'][c]['htlcs']) == 0)
 
+    # This sends money to a directly connected peer
     def pay(self, dst, amt, label=None):
         if not label:
             label = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(20))
@@ -839,6 +841,18 @@ class LightningNode(object):
         result = self.rpc.waitsendpay(rhash)
         assert(result.get('status') == 'complete')
 
+    # This helper sends all money to a peer until even 1 msat can't get through.
+    def drain(self, peer):
+        total = 0
+        msat = 16**9
+        while msat != 0:
+            try:
+                self.pay(peer, msat)
+                total += msat
+            except RpcError:
+                msat //= 2
+        return total
+
     # Note: this feeds through the smoother in update_feerate, so changing
     # it on a running daemon may not give expected result!
     def set_feerates(self, feerates, wait_for_effect=True):
@@ -868,6 +882,15 @@ class LightningNode(object):
         if wait_for_effect:
             wait_for(lambda: self.daemon.rpcproxy.mock_counts['estimatesmartfee'] >= 3)
 
+    # force new feerates by restarting and thus skipping slow smoothed process
+    # Note: testnode must be created with: opts={'may_reconnect': True}
+    def force_feerates(self, rate):
+        assert(self.may_reconnect)
+        self.set_feerates([rate] * 3, False)
+        self.restart()
+        self.daemon.wait_for_log('peer_out WIRE_UPDATE_FEE')
+        assert(self.rpc.feerates('perkw')['perkw']['normal'] == rate)
+
     def wait_for_onchaind_broadcast(self, name, resolve=None):
         """Wait for onchaind to drop tx name to resolve (if any)"""
         if resolve:
diff --git a/tests/test_pay.py b/tests/test_pay.py
index b187b405f..7ad1792d4 100644
--- a/tests/test_pay.py
+++ b/tests/test_pay.py
@@ -2275,34 +2275,20 @@ def test_lockup_drain(node_factory, bitcoind):
     l1, l2 = node_factory.line_graph(2, opts={'may_reconnect': True})
 
     # l1 sends all the money to l2 until even 1 msat can't get through.
-    total = 0
-    msat = 16**9
-    while msat != 0:
-        try:
-            l1.pay(l2, msat)
-            print("l1->l2: {}".format(msat))
-            total += msat
-        except RpcError:
-            msat //= 2
+    total = l1.drain(l2)
 
     # Even if feerate now increases 2x (30000), l2 should be able to send
     # non-dust HTLC to l1.
-    l1.set_feerates([30000] * 3, False)
-    # Restart forces fast fee adjustment (otherwise it's smoothed and takes
-    # a very long time!).
-    l1.restart()
-    wait_for(lambda: only_one(l1.rpc.listpeers()['peers'])['connected'])
-    assert(l1.rpc.feerates('perkw')['perkw']['normal'] == 30000)
-
+    l1.force_feerates(30000)
     l2.pay(l1, total // 2)
 
-    # But if feerate increase just a little more, l2 should not be able to send
-    # non-dust HTLC to l1
-    l1.set_feerates([30002] * 3, False)  # TODO: why does 30001 fail, off by one in C code?
-    l1.restart()
-    wait_for(lambda: only_one(l1.rpc.listpeers()['peers'])['connected'])
-    assert(l1.rpc.feerates('perkw')['perkw']['normal'] == 30002)
+    # reset fees and send all back again
+    l1.force_feerates(15000)
+    l1.drain(l2)
 
+    # But if feerate increase just a little more, l2 should not be able to send
+    # non-fust HTLC to l1
+    l1.force_feerates(30002)  # TODO: Why does 30001 fail? off by one in C code?
     with pytest.raises(RpcError, match=r".*Capacity exceeded.*"):
         l2.pay(l1, total // 2)