From 0d75f0a2c77b05b247adac3a08371203d12ac1d6 Mon Sep 17 00:00:00 2001
From: Antoine Poinsot <darosior@protonmail.com>
Date: Tue, 28 Mar 2023 15:48:42 +0200
Subject: [PATCH] tests: handle multiple recovery paths in MultiSigner

---
 tests/fixtures.py              |  4 ++--
 tests/test_framework/signer.py | 35 +++++++++++++++++++++++-----------
 tests/test_misc.py             |  2 +-
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/tests/fixtures.py b/tests/fixtures.py
index 0ef108a5..4122e13b 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -163,11 +163,11 @@ def lianad_multisig(bitcoind, directory):
     bitcoind_cookie = os.path.join(bitcoind.bitcoin_dir, "regtest", ".cookie")
 
     # A 3-of-4 that degrades into a 2-of-5 after 10 blocks
-    signer = MultiSigner(4, 5)
     csv_value = 10
+    signer = MultiSigner(4, {csv_value: 5})
     prim_multi, recov_multi = (
         multi_expression(3, signer.prim_hds),
-        multi_expression(2, signer.recov_hds),
+        multi_expression(2, signer.recov_hds[csv_value]),
     )
     main_desc = Descriptor.from_str(
         f"wsh(or_d({prim_multi},and_v(v:{recov_multi},older({csv_value}))))"
diff --git a/tests/test_framework/signer.py b/tests/test_framework/signer.py
index e1af3772..9e0f86cc 100644
--- a/tests/test_framework/signer.py
+++ b/tests/test_framework/signer.py
@@ -59,6 +59,8 @@ def sign_psbt(psbt, hds):
 
 
 class SingleSigner:
+    """Assumes a simple 1-primary path 1-recovery path Liana descriptor."""
+
     def __init__(self):
         self.primary_hd = BIP32.from_seed(os.urandom(32), network="test")
         self.recovery_hd = BIP32.from_seed(os.urandom(32), network="test")
@@ -76,20 +78,31 @@ class SingleSigner:
 
 
 class MultiSigner:
-    def __init__(
-        self, primary_hds_count, recovery_hds_count
-    ):
+    """A signer that has multiple keys and may have multiple recovery path."""
+
+    def __init__(self, primary_hds_count, recovery_hds_counts):
         self.prim_hds = [
             BIP32.from_seed(os.urandom(32), network="test")
             for _ in range(primary_hds_count)
         ]
-        self.recov_hds = [
-            BIP32.from_seed(os.urandom(32), network="test")
-            for _ in range(recovery_hds_count)
-        ]
+        self.recov_hds = {}
+        for timelock, count in recovery_hds_counts.items():
+            self.recov_hds[timelock] = [
+                BIP32.from_seed(os.urandom(32), network="test") for _ in range(count)
+            ]
 
-    def sign_psbt(self, psbt, key_indices, recovery=False):
-        """Sign a transaction with the keys at the specified indices."""
-        hds = self.recov_hds if recovery else self.prim_hds
-        hds = [hds[i] for i in key_indices]
+    def sign_psbt(self, psbt, key_indices):
+        """Sign a transaction with the keys at the specified indices.
+
+        The key indices may be specified as a mapping from timelock value to list of
+        indices to sign with the keys of a specific recovery path.
+        """
+        if isinstance(key_indices, dict):
+            hds = [
+                self.recov_hds[timelock][i]
+                for timelock, indices in key_indices.items()
+                for i in indices
+            ]
+        else:
+            hds = [self.prim_hds[i] for i in key_indices]
         return sign_psbt(psbt, hds)
diff --git a/tests/test_misc.py b/tests/test_misc.py
index 483c3f3c..5222f713 100644
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -70,7 +70,7 @@ def test_multisig(lianad_multisig, bitcoind):
     res = lianad_multisig.rpc.createrecovery(bitcoind.rpc.getnewaddress(), 2)
     reco_psbt = PSBT.from_base64(res["psbt"])
     txid = reco_psbt.tx.txid().hex()
-    signed_psbt = lianad_multisig.signer.sign_psbt(reco_psbt, [1, 4], recovery=True)
+    signed_psbt = lianad_multisig.signer.sign_psbt(reco_psbt, {10: [1, 4]})
     lianad_multisig.rpc.updatespend(signed_psbt.to_base64())
     lianad_multisig.rpc.broadcastspend(txid)