path: root/btc-wire/src/bin/test.rs

use core::panic;
use std::{collections::HashSet, iter::repeat_with, panic::AssertUnwindSafe};

use bitcoincore_rpc::{
    bitcoin::{Address, Amount, Txid},
    json::GetTransactionResultDetailCategory as Category,
    jsonrpc::{
        error::RpcError,
        serde_json::{json, Value},
    },
    Client, RpcApi,
};
use btc_wire::{
    rpc_patch::RpcErrorCode,
    rpc_utils::{
        common_rpc, default_data_dir, dirty_guess_network, wallet_rpc, Network, CLIENT, WIRE,
    },
    test::rand_key,
    BounceErr, ClientExtended,
};
use owo_colors::OwoColorize;

const RESERVE: &str = "reserve";

/// Instrumentation test
pub fn main() {
    let test_amount = Amount::from_sat(1500);
    let data_dir = default_data_dir();
    // Network check
    let network = dirty_guess_network(&data_dir);
    match network {
        Network::MainNet => {
            panic!("Do not run tests on the mainnet, you are going to loose money")
        }
        Network::TestNet => println!("{}", "Running on testnet, slow network mining".yellow()),
        Network::RegTest => println!("Running on regtest, fast manual mining"),
    }

    // Wallet check
    {
        let existing_wallets: HashSet<String> =
            std::fs::read_dir(data_dir.join(network.dir()).join("wallets"))
                .unwrap()
                .filter_map(|it| it.ok())
                .map(|it| it.file_name().to_string_lossy().to_string())
                .collect();

        let rpc = common_rpc(&data_dir, network).expect("Failed to open common client");
        if !existing_wallets.contains(CLIENT)
            || !existing_wallets.contains(WIRE)
            || !existing_wallets.contains(RESERVE)
        {
            println!("Generate tests wallets");
            // Create wallets
            rpc.create_wallet(WIRE, None, None, None, None).ok();
            rpc.create_wallet(CLIENT, None, None, None, None).ok();
            rpc.create_wallet(RESERVE, None, None, None, None).ok();
        }

        // Load wallets

        rpc.load_wallet(WIRE).ok();
        rpc.load_wallet(CLIENT).ok();
        rpc.load_wallet(RESERVE).ok();
    }

    // Client initialization
    let client_rpc = wallet_rpc(&data_dir, network, CLIENT);
    let wire_rpc = wallet_rpc(&data_dir, network, WIRE);
    let reserve_rpc = wallet_rpc(&data_dir, network, RESERVE);
    let client_addr = client_rpc.get_new_address(None, None).unwrap();
    let wire_addr = wire_rpc.get_new_address(None, None).unwrap();
    let reserve_addr = reserve_rpc.get_new_address(None, None).unwrap();

    let next_block = || {
        match network {
            Network::RegTest => {
                // Manually mine a block
                reserve_rpc.generate_to_address(1, &reserve_addr).unwrap();
            }
            _ => {
                // Wait for next network block
                reserve_rpc.wait_for_new_block(0).ok();
            }
        }
    };

    let wait_for_tx = |rpc: &Client, txs: &[Txid]| {
        let mut count = 0;
        while txs
            .iter()
            .any(|id| rpc.get_transaction(id, None).unwrap().info.confirmations <= 0)
        {
            next_block();
            if count > 3 {
                panic!("Transaction no sended after 4 blocks");
            }
            count += 1;
        }
    };

    // Balance check
    {
        // Transfer all wire money to client
        let wire_balance = wire_rpc.get_balance(None, None).unwrap();
        wire_rpc
            .send_to_address(
                &client_addr,
                wire_balance,
                None,
                None,
                Some(true),
                None,
                None,
                None,
            )
            .ok();
        // Transfer all wire money to client
        let reserve_balance = reserve_rpc.get_balance(None, None).unwrap();
        reserve_rpc
            .send_to_address(
                &client_addr,
                reserve_balance,
                None,
                None,
                Some(true),
                None,
                None,
                None,
            )
            .ok();
        next_block();

        let balance = client_rpc.get_balance(None, None).unwrap();
        let min_balance = test_amount * 3;
        if balance < min_balance {
            println!(
                "{}",
                format_args!(
                    "Client wallet have only {}, {} are required to perform the test",
                    balance, min_balance
                )
                .yellow()
            );
            match network {
                Network::MainNet | Network::TestNet => {
                    if network == Network::MainNet {
                        println!("Send coins to this address: {}", client_addr);
                    } else {
                        println!("Request coins from a faucet such as https://bitcoinfaucet.uo1.net/send.php to this address: {}", client_addr);
                    }
                    println!("Waiting for the transaction...");
                    while client_rpc.get_balance(None, None).unwrap() < min_balance {
                        client_rpc.wait_for_new_block(0).ok();
                    }
                }
                Network::RegTest => {
                    println!("Add 50B to client wallet");
                    client_rpc
                        .generate_to_address(
                            101, /* Need 100 blocks to validate */
                            &client_addr,
                        )
                        .unwrap();
                }
            }
        }

        println!(
            "Initial state:\n{} {}\n{} {}",
            WIRE,
            wire_rpc.get_balance(None, None).unwrap(),
            CLIENT,
            client_rpc.get_balance(None, None).unwrap()
        );
    }

    let mut runner = TestRunner::new();

    runner.test("OpReturn metadata", || {
        // Send metadata
        let msg = "J'aime le chocolat".as_bytes();
        let id = client_rpc
            .send_op_return(&wire_addr, test_amount, msg)
            .unwrap();
        // Check in mempool
        assert!(
            tx_exist(&client_rpc, &id, 0, Category::Send).unwrap(),
            "Not in mempool"
        );
        // Check mined
        wait_for_tx(&client_rpc, &[id]);
        assert!(
            tx_exist(&wire_rpc, &id, 1, Category::Receive).unwrap(),
            "Not mined"
        );
        // Check extract
        let (_, extracted) = wire_rpc.get_tx_op_return(&id).unwrap();
        assert_eq!(msg, extracted, "Corrupted metadata");
    });
    runner.test("SegWit metadata", || {
        // Send metadata
        let key = rand_key();
        let id = client_rpc
            .send_segwit_key(&wire_addr, test_amount, &key)
            .unwrap();
        // Check in mempool
        assert!(
            tx_exist(&client_rpc, &id, 0, Category::Send).unwrap(),
            "Not in mempool"
        );
        // Check mined
        wait_for_tx(&client_rpc, &[id]);
        assert!(
            tx_exist(&wire_rpc, &id, 1, Category::Receive).unwrap(),
            "Not mined"
        );
        // Check extract
        let (_, extracted) = wire_rpc.get_tx_segwit_key(&id).unwrap();
        assert_eq!(key, extracted, "Corrupted metadata");
    });

    let bounce_fee = Amount::from_sat(300);
    runner.test("Bounce simple", || {
        let before = client_rpc.get_balance(None, None).unwrap();
        let send_id = client_rpc
            .send_to_address(&wire_addr, test_amount, None, None, None, None, None, None)
            .unwrap();
        wait_for_tx(&client_rpc, &[send_id]);
        let bounce_id = wire_rpc.bounce(&send_id, bounce_fee).unwrap();
        wait_for_tx(&wire_rpc, &[bounce_id]);
        let bounce_tx_fee = wire_rpc.get_transaction(&bounce_id, None).unwrap().details[0]
            .fee
            .unwrap()
            .abs()
            .to_unsigned()
            .unwrap();
        let send_tx_fee = client_rpc.get_transaction(&send_id, None).unwrap().details[0]
            .fee
            .unwrap()
            .abs()
            .to_unsigned()
            .unwrap();
        let after = client_rpc.get_balance(None, None).unwrap();
        assert!(before >= after);
        assert_eq!(before - after, bounce_tx_fee + bounce_fee + send_tx_fee);
    });
    runner.test("Bounce minimal amount", || {
        let send_id = client_rpc
            .send_to_address(
                &wire_addr,
                Amount::from_sat(294),
                None,
                None,
                None,
                None,
                None,
                None,
            )
            .unwrap();
        wait_for_tx(&client_rpc, &[send_id]);
        assert!(match wire_rpc.bounce(&send_id, bounce_fee) {
            Ok(_) => false,
            Err(err) => match err {
                BounceErr::AmountLessThanFee => true,
                _ => false,
            },
        });
    });
    runner.test("Bounce too small amount", || {
        let send_id = client_rpc
            .send_to_address(
                &wire_addr,
                Amount::from_sat(294) + bounce_fee,
                None,
                None,
                None,
                None,
                None,
                None,
            )
            .unwrap();
        wait_for_tx(&client_rpc, &[send_id]);

        assert!(match wire_rpc.bounce(&send_id, bounce_fee) {
            Ok(_) => false,
            Err(err) => match err {
                BounceErr::RPC(bitcoincore_rpc::Error::JsonRpc(
                    bitcoincore_rpc::jsonrpc::Error::Rpc(RpcError { code, .. }),
                )) => code == RpcErrorCode::RpcWalletInsufficientFunds as i32,
                _ => false,
            },
        });
    });
    runner.test("Bounce complex", || {
        // Generate 6 new addresses
        let addresses: Vec<Address> =
            repeat_with(|| client_rpc.get_new_address(None, None).unwrap())
                .take(6)
                .collect();
        // Send transaction to self with 1, 2 and 3 outputs
        let txs: Vec<Txid> = [&addresses[0..1], &addresses[1..3], &addresses[3..]]
            .into_iter()
            .map(|addresses| {
                let hex: String = client_rpc
                    .call(
                        "createrawtransaction",
                        &[
                            Value::Array(vec![]),
                            Value::Array(
                                addresses
                                    .iter()
                                    .map(|addr| json!({&addr.to_string(): test_amount.as_btc()}))
                                    .collect(),
                            ),
                        ],
                    )
                    .unwrap();
                let funded = client_rpc.fund_raw_transaction(hex, None, None).unwrap();
                let signed = client_rpc
                    .sign_raw_transaction_with_wallet(&funded.hex, None, None)
                    .unwrap();
                client_rpc.send_raw_transaction(&signed.hex).unwrap()
            })
            .collect();
        wait_for_tx(&client_rpc, txs.as_slice());
        let before = client_rpc.get_balance(None, None).unwrap();
        // Send a transaction with multiple input from multiple transaction of different outputs len
        let hex: String = client_rpc
            .call(
                "createrawtransaction",
                &[
                    Value::Array(
                        txs.into_iter()
                            .enumerate()
                            .map(|(i, tx)| {
                                json!({
                                    "txid": tx.to_string(),
                                    "vout": i,
                                })
                            })
                            .collect(),
                    ),
                    json!([
                        {&wire_addr.to_string(): (test_amount*3).as_btc()}
                    ]),
                ],
            )
            .unwrap();
        let funded = client_rpc.fund_raw_transaction(hex, None, None).unwrap();
        let signed = client_rpc
            .sign_raw_transaction_with_wallet(&funded.hex, None, None)
            .unwrap();
        let send_id = client_rpc.send_raw_transaction(&signed.hex).unwrap();
        wait_for_tx(&client_rpc, &[send_id]);
        let bounce_id = wire_rpc.bounce(&send_id, bounce_fee).unwrap();
        wait_for_tx(&wire_rpc, &[bounce_id]);
        let after = client_rpc.get_balance(None, None).unwrap();
        let bounce_tx_fee = wire_rpc.get_transaction(&bounce_id, None).unwrap().details[0]
            .fee
            .unwrap()
            .abs()
            .to_unsigned()
            .unwrap();
        let send_tx_fee = client_rpc.get_transaction(&send_id, None).unwrap().details[0]
            .fee
            .unwrap()
            .abs()
            .to_unsigned()
            .unwrap();
        assert!(before >= after);
        assert_eq!(before - after, bounce_tx_fee + bounce_fee + send_tx_fee);
    });

    runner.conclude();
}

/// Check a specific transaction exist in a wallet historic
fn tx_exist(
    rpc: &Client,
    id: &Txid,
    min_confirmation: i32,
    detail: Category,
) -> bitcoincore_rpc::Result<bool> {
    let txs = rpc.list_transactions(None, None, None, None).unwrap();
    let found = txs.into_iter().any(|tx| {
        tx.detail.category == detail
            && tx.info.confirmations >= min_confirmation
            && tx.info.txid == *id
    });
    Ok(found)
}

/// Run test track success and errors
struct TestRunner {
    nb_ok: usize,
    nb_err: usize,
}

impl TestRunner {
    fn new() -> Self {
        Self {
            nb_err: 0,
            nb_ok: 0,
        }
    }

    fn test(&mut self, name: &str, test: impl FnOnce()) {
        println!("{}", name.cyan());

        let result = std::panic::catch_unwind(AssertUnwindSafe(test));
        if result.is_ok() {
            println!("{}", "OK".green());
            self.nb_ok += 1;
        } else {
            dbg!(&result);
            println!("{}", "ERR".red());
            self.nb_err += 1;
        }
    }

    /// Wait for tests completion and print results
    fn conclude(self) {
        println!(
            "Result for {} tests: {} ok and {} err",
            self.nb_ok + self.nb_err,
            self.nb_ok,
            self.nb_err
        );
    }
}