Merge pull request #2966 from G8XSU/2647-distribute

[rust-lightning] / lightning / src / util / test_utils.rs

diff --git a/lightning/src/util/test_utils.rs b/lightning/src/util/test_utils.rs
index 2d00844772aca60eb841b2eb5bde57f7dcae8556..4d17f8fd17ddcd9a919acb4a4abddd9aac673c17 100644 (file)
--- a/lightning/src/util/test_utils.rs
+++ b/lightning/src/util/test_utils.rs
@@ -79,6 +79,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use bitcoin::psbt::Psbt;
 use bitcoin::Sequence;
 
+use super::test_channel_signer::SignerOp;
+
 pub fn pubkey(byte: u8) -> PublicKey {
        let secp_ctx = Secp256k1::new();
        PublicKey::from_secret_key(&secp_ctx, &privkey(byte))
@@ -1222,6 +1224,7 @@ pub struct TestKeysInterface {
        enforcement_states: Mutex<HashMap<[u8;32], Arc<Mutex<EnforcementState>>>>,
        expectations: Mutex<Option<VecDeque<OnGetShutdownScriptpubkey>>>,
        pub unavailable_signers: Mutex<HashSet<[u8; 32]>>,
+       pub unavailable_signers_ops: Mutex<HashMap<[u8; 32], HashSet<SignerOp>>>,
 }
 
 impl EntropySource for TestKeysInterface {
@@ -1280,9 +1283,11 @@ impl SignerProvider for TestKeysInterface {
        fn derive_channel_signer(&self, channel_value_satoshis: u64, channel_keys_id: [u8; 32]) -> TestChannelSigner {
                let keys = self.backing.derive_channel_signer(channel_value_satoshis, channel_keys_id);
                let state = self.make_enforcement_state_cell(keys.commitment_seed);
-               let signer = TestChannelSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check);
-               if self.unavailable_signers.lock().unwrap().contains(&channel_keys_id) {
-                       signer.set_available(false);
+               let mut signer = TestChannelSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check);
+               if let Some(ops) = self.unavailable_signers_ops.lock().unwrap().get(&channel_keys_id) {
+                       for &op in ops {
+                               signer.disable_op(op);
+                       }
                }
                signer
        }
@@ -1323,6 +1328,7 @@ impl TestKeysInterface {
                        enforcement_states: Mutex::new(new_hash_map()),
                        expectations: Mutex::new(None),
                        unavailable_signers: Mutex::new(new_hash_set()),
+                       unavailable_signers_ops: Mutex::new(new_hash_map()),
                }
        }