Make mining more flexible in bdk_testenv

crates/testenv/src/lib.rs

@@ -3,9 +3,16 @@
 pub mod utils;
 
 use anyhow::Context;
+use bdk_chain::bitcoin::{
+    block::Header, hash_types::TxMerkleNode, hex::FromHex, script::PushBytesBuf, transaction,
+    Address, Amount, Block, BlockHash, ScriptBuf, Transaction, TxIn, TxOut, Txid,
+};
 use bdk_chain::CheckPoint;
-use bitcoin::{address::NetworkChecked, Address, Amount, BlockHash, Txid};
-use std::time::Duration;
+use bitcoin::address::NetworkChecked;
+use bitcoin::hex::HexToBytesError;
+use core::time::Duration;
+use electrsd::corepc_node::mtype::GetBlockTemplate;
+use electrsd::corepc_node::{TemplateRequest, TemplateRules};
 
 pub use electrsd;
 pub use electrsd::corepc_client;
@@ -45,6 +52,32 @@ impl Default for Config<'_> {
     }
 }
 
+/// Parameters for [`TestEnv::mine_block`].
+#[non_exhaustive]
+#[derive(Default)]
+pub struct MineParams {
+    /// If `true`, the block will be empty (no mempool transactions).
+    pub empty: bool,
+    /// Set a custom block timestamp. Defaults to `max(min_time, now)`.
+    pub time: Option<u32>,
+    /// Set a custom coinbase output script. Defaults to `OP_TRUE`.
+    pub coinbase_address: Option<ScriptBuf>,
+}
+
+impl MineParams {
+    fn address_or_anyone_can_spend(&self) -> ScriptBuf {
+        use bdk_chain::bitcoin::opcodes::OP_TRUE;
+        self.coinbase_address
+            .clone()
+            // OP_TRUE (anyone can spend)
+            .unwrap_or_else(|| {
+                bdk_chain::bitcoin::script::Builder::new()
+                    .push_opcode(OP_TRUE)
+                    .into_script()
+            })
+    }
+}
+
 impl TestEnv {
     /// Construct a new [`TestEnv`] instance with the default configuration used by BDK.
     pub fn new() -> anyhow::Result<Self> {
@@ -119,52 +152,135 @@ impl TestEnv {
         Ok(block_hashes)
     }
 
+    /// Get a block template from the node.
+    pub fn get_block_template(&self) -> anyhow::Result<GetBlockTemplate> {
+        Ok(self
+            .bitcoind
+            .client
+            .get_block_template(&TemplateRequest {
+                rules: vec![
+                    TemplateRules::Segwit,
+                    TemplateRules::Taproot,
+                    TemplateRules::Csv,
+                ],
+            })?
+            .into_model()?)
+    }
+
     /// Mine a block that is guaranteed to be empty even with transactions in the mempool.
     #[cfg(feature = "std")]
     pub fn mine_empty_block(&self) -> anyhow::Result<(usize, BlockHash)> {
-        use bitcoin::secp256k1::rand::random;
-        use bitcoin::{
-            block::Header, hashes::Hash, transaction, Block, ScriptBuf, ScriptHash, Transaction,
-            TxIn, TxMerkleNode, TxOut,
+        self.mine_block(MineParams {
+            empty: true,
+            ..Default::default()
+        })
+    }
+
+    /// Get the minimum valid timestamp for the next block.
+    pub fn min_time_for_next_block(&self) -> anyhow::Result<u32> {
+        Ok(self.get_block_template()?.min_time)
+    }
+
+    /// Mine a single block with the given [`MineParams`].
+    pub fn mine_block(&self, params: MineParams) -> anyhow::Result<(usize, BlockHash)> {
+        let bt = self.get_block_template()?;
+
+        // BIP34 requires the height to be the first item in coinbase scriptSig.
+        // Bitcoin Core validates by checking if scriptSig STARTS with the expected
+        // encoding (using minimal opcodes like OP_1 for height 1).
+        // The scriptSig must also be 2-100 bytes total.
+        fn build_coinbase_scriptsig(bt: &GetBlockTemplate, pad: bool) -> ScriptBuf {
+            let mut builder = bdk_chain::bitcoin::script::Builder::new().push_int(bt.height as i64);
+            if pad {
+                builder = builder.push_opcode(bdk_chain::bitcoin::opcodes::OP_0);
+            }
+            for v in bt.coinbase_aux.values() {
+                let bytes = Vec::<u8>::from_hex(v).expect("must be valid hex");
+                let bytes_buf = PushBytesBuf::try_from(bytes).expect("must be valid bytes");
+                builder = builder.push_slice(bytes_buf);
+            }
+            builder.into_script()
+        }
+        let coinbase_scriptsig = {
+            let mut script = build_coinbase_scriptsig(&bt, false);
+            // Ensure scriptSig is at least 2 bytes (pad with OP_0 if needed)
+            if script.len() < 2 {
+                script = build_coinbase_scriptsig(&bt, true);
+            };
+            script
         };
-        use corepc_node::{TemplateRequest, TemplateRules};
-        let request = TemplateRequest {
-            rules: vec![TemplateRules::Segwit],
+
+        let coinbase_outputs = if params.empty {
+            let tx_fees: Amount = bt
+                .transactions
+                .iter()
+                .map(|tx| tx.fee.to_unsigned().expect("fee must be positive"))
+                .sum();
+            let value = bt
+                .coinbase_value
+                .to_unsigned()
+                .expect("coinbase_value must be positive")
+                - tx_fees;
+            vec![TxOut {
+                value,
+                script_pubkey: params.address_or_anyone_can_spend(),
+            }]
+        } else {
+            core::iter::once(TxOut {
+                value: bt
+                    .coinbase_value
+                    .to_unsigned()
+                    .expect("coinbase_value must be positive"),
+                script_pubkey: params.address_or_anyone_can_spend(),
+            })
+            .chain(
+                bt.default_witness_commitment
+                    .as_ref()
+                    .map(|s| -> Result<_, HexToBytesError> {
+                        Ok(TxOut {
+                            value: Amount::ZERO,
+                            script_pubkey: ScriptBuf::from_hex(s)?,
+                        })
+                    })
+                    .transpose()?,
+            )
+            .collect()
         };
-        let bt = self
-            .bitcoind
-            .client
-            .get_block_template(&request)?
-            .into_model()?;
 
-        let txdata = vec![Transaction {
+        let coinbase_tx = Transaction {
             version: transaction::Version::ONE,
             lock_time: bdk_chain::bitcoin::absolute::LockTime::from_height(0)?,
             input: vec![TxIn {
                 previous_output: bdk_chain::bitcoin::OutPoint::default(),
-                script_sig: ScriptBuf::builder()
-                    .push_int(bt.height as _)
-                    // random number so that re-mining creates unique block
-                    .push_int(random())
-                    .into_script(),
+                script_sig: coinbase_scriptsig,
                 sequence: bdk_chain::bitcoin::Sequence::default(),
                 witness: bdk_chain::bitcoin::Witness::new(),
             }],
-            output: vec![TxOut {
-                value: Amount::ZERO,
-                script_pubkey: ScriptBuf::new_p2sh(&ScriptHash::all_zeros()),
-            }],
-        }];
+            output: coinbase_outputs,
+        };
+
+        let txdata = if params.empty {
+            vec![coinbase_tx]
+        } else {
+            core::iter::once(coinbase_tx)
+                .chain(bt.transactions.iter().map(|tx| tx.data.clone()))
+                .collect()
+        };
 
         let mut block = Block {
             header: Header {
                 version: bt.version,
                 prev_blockhash: bt.previous_block_hash,
-                merkle_root: TxMerkleNode::all_zeros(),
-                time: Ord::max(
+                merkle_root: TxMerkleNode::from_raw_hash(
+                    bdk_chain::bitcoin::merkle_tree::calculate_root(
+                        txdata.iter().map(|tx| tx.compute_txid().to_raw_hash()),
+                    )
+                    .expect("must have atleast one tx"),
+                ),
+                time: params.time.unwrap_or(Ord::max(
                     bt.min_time,
                     std::time::UNIX_EPOCH.elapsed()?.as_secs() as u32,
-                ),
+                )),
                 bits: bt.bits,
                 nonce: 0,
             },
@@ -173,16 +289,18 @@ impl TestEnv {
 
         block.header.merkle_root = block.compute_merkle_root().expect("must compute");
 
+        // Mine!
+        let target = block.header.target();
         for nonce in 0..=u32::MAX {
             block.header.nonce = nonce;
-            if block.header.target().is_met_by(block.block_hash()) {
-                break;
+            let blockhash = block.block_hash();
+            if target.is_met_by(blockhash) {
+                self.rpc_client().submit_block(&block)?;
+                return Ok((bt.height as usize, blockhash));
             }
         }
 
-        self.bitcoind.client.submit_block(&block)?;
-
-        Ok((bt.height as usize, block.block_hash()))
+        Err(anyhow::anyhow!("Cannot find nonce that meets the target"))
     }
 
     /// This method waits for the Electrum notification indicating that a new block has been mined.
@@ -318,9 +436,12 @@ impl TestEnv {
 #[cfg(test)]
 #[cfg_attr(coverage_nightly, coverage(off))]
 mod test {
-    use crate::TestEnv;
+    use crate::{MineParams, TestEnv};
+    use bdk_chain::bitcoin::opcodes::OP_TRUE;
+    use bdk_chain::bitcoin::Amount;
     use core::time::Duration;
     use electrsd::corepc_node::anyhow::Result;
+    use std::collections::BTreeSet;
 
     /// This checks that reorgs initiated by `bitcoind` is detected by our `electrsd` instance.
     #[test]
@@ -355,4 +476,77 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_mine_block() -> Result<()> {
+        let anyone_can_spend = bdk_chain::bitcoin::script::Builder::new()
+            .push_opcode(OP_TRUE)
+            .into_script();
+
+        let env = TestEnv::new()?;
+
+        // So we can spend.
+        let addr = env
+            .rpc_client()
+            .get_new_address(None, None)?
+            .address()?
+            .assume_checked();
+        env.mine_blocks(100, Some(addr.clone()))?;
+
+        // Try mining a block with custom time.
+        let custom_time = env.min_time_for_next_block()? + 100;
+        let (_a_height, a_hash) = env.mine_block(MineParams {
+            empty: false,
+            time: Some(custom_time),
+            coinbase_address: None,
+        })?;
+        let a_block = env.rpc_client().get_block(a_hash)?;
+        assert_eq!(a_block.header.time, custom_time);
+        assert_eq!(
+            a_block.txdata[0].output[0].script_pubkey, anyone_can_spend,
+            "Subsidy address must be anyone_can_spend"
+        );
+
+        // Now try mining with min time & some txs.
+        let txid1 = env.send(&addr, Amount::from_sat(100_000))?;
+        let txid2 = env.send(&addr, Amount::from_sat(200_000))?;
+        let txid3 = env.send(&addr, Amount::from_sat(300_000))?;
+        let min_time = env.min_time_for_next_block()?;
+        let (_b_height, b_hash) = env.mine_block(MineParams {
+            empty: false,
+            time: Some(min_time),
+            coinbase_address: None,
+        })?;
+        let b_block = env.rpc_client().get_block(b_hash)?;
+        assert_eq!(b_block.header.time, min_time);
+        assert_eq!(
+            a_block.txdata[0].output[0].script_pubkey, anyone_can_spend,
+            "Subsidy address must be anyone_can_spend"
+        );
+        assert_eq!(
+            b_block
+                .txdata
+                .iter()
+                .skip(1) // ignore coinbase
+                .map(|tx| tx.compute_txid())
+                .collect::<BTreeSet<_>>(),
+            [txid1, txid2, txid3].into_iter().collect(),
+            "Must have all txs"
+        );
+
+        // Custom subsidy address.
+        let (_c_height, c_hash) = env.mine_block(MineParams {
+            empty: false,
+            time: None,
+            coinbase_address: Some(addr.script_pubkey()),
+        })?;
+        let c_block = env.rpc_client().get_block(c_hash)?;
+        assert_eq!(
+            c_block.txdata[0].output[0].script_pubkey,
+            addr.script_pubkey(),
+            "Custom address works"
+        );
+
+        Ok(())
+    }
 }