From fc877bd89e0be387679d9b8aee8eda5571940117 Mon Sep 17 00:00:00 2001
From: Runchao Han <me@runchao.rocks>
Date: Fri, 30 Aug 2024 14:57:48 +1000
Subject: [PATCH] init

---
 contracts/btc-staking/src/contract.rs | 28 ++++++++++++++-------------
 packages/btcstaking/src/tx_verify.rs  | 14 +++++++++-----
 packages/test-utils/src/lib.rs        | 14 +++++++-------
 3 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/contracts/btc-staking/src/contract.rs b/contracts/btc-staking/src/contract.rs
index 9978752a..6a1ff385 100644
--- a/contracts/btc-staking/src/contract.rs
+++ b/contracts/btc-staking/src/contract.rs
@@ -277,30 +277,26 @@ fn handle_end_block(
 pub(crate) mod tests {
     use super::*;
 
-    use hex::ToHex;
-
+    use babylon_apis::btc_staking_api::{
+        ActiveBtcDelegation, BtcUndelegationInfo, CovenantAdaptorSignatures,
+        FinalityProviderDescription, NewFinalityProvider, ProofOfPossessionBtc,
+    };
+    use babylon_apis::finality_api::PubRandCommit;
+    use babylon_proto::babylon::btcstaking::v1::BtcDelegation;
     use cosmwasm_std::{
         from_json,
         testing::{message_info, mock_dependencies, mock_env},
         Binary, Decimal,
     };
     use cw_controllers::AdminResponse;
-
-    use babylon_apis::btc_staking_api::{
-        ActiveBtcDelegation, BtcUndelegationInfo, CovenantAdaptorSignatures,
-        FinalityProviderDescription, NewFinalityProvider, ProofOfPossessionBtc,
-    };
-    use babylon_apis::finality_api::PubRandCommit;
-
-    use test_utils::{get_btc_delegation_and_params, get_pub_rand_commit};
+    use hex::ToHex;
+    use test_utils::{get_btc_delegation, get_pub_rand_commit};
 
     pub(crate) const CREATOR: &str = "creator";
     pub(crate) const INIT_ADMIN: &str = "initial_admin";
     const NEW_ADMIN: &str = "new_admin";
 
-    /// Build an active BTC delegation from a BTC delegation
-    pub(crate) fn get_active_btc_delegation() -> ActiveBtcDelegation {
-        let (del, _) = get_btc_delegation_and_params();
+    fn new_active_btc_delegation(del: BtcDelegation) -> ActiveBtcDelegation {
         let btc_undelegation = del.btc_undelegation.unwrap();
 
         ActiveBtcDelegation {
@@ -347,6 +343,12 @@ pub(crate) mod tests {
         }
     }
 
+    /// Build an active BTC delegation from a BTC delegation
+    pub(crate) fn get_active_btc_delegation() -> ActiveBtcDelegation {
+        let del = get_btc_delegation();
+        new_active_btc_delegation(del)
+    }
+
     // Build a derived active BTC delegation from the base (from testdata) BTC delegation
     pub(crate) fn get_derived_btc_delegation(del_id: u8, fp_ids: &[u8]) -> ActiveBtcDelegation {
         assert!(
diff --git a/packages/btcstaking/src/tx_verify.rs b/packages/btcstaking/src/tx_verify.rs
index 672e38c8..569169b9 100644
--- a/packages/btcstaking/src/tx_verify.rs
+++ b/packages/btcstaking/src/tx_verify.rs
@@ -204,11 +204,12 @@ mod tests {
     use bitcoin::consensus::deserialize;
     use bitcoin::secp256k1::schnorr::Signature;
     use bitcoin::{Transaction, XOnlyPublicKey};
-    use test_utils::get_btc_delegation_and_params;
+    use test_utils::{get_btc_delegation, get_params};
 
     #[test]
     fn test_check_transactions() {
-        let (btc_del, params) = get_btc_delegation_and_params();
+        let btc_del = get_btc_delegation();
+        let params = get_params();
 
         let staking_tx: Transaction = deserialize(&btc_del.staking_tx).unwrap();
         let slashing_tx: Transaction = deserialize(&btc_del.slashing_tx).unwrap();
@@ -239,7 +240,8 @@ mod tests {
 
     #[test]
     fn test_verify_unbonding_tx_schnorr_sig() {
-        let (btc_del, params) = get_btc_delegation_and_params();
+        let btc_del = get_btc_delegation();
+        let params = get_params();
 
         let staking_tx: Transaction = deserialize(&btc_del.staking_tx).unwrap();
         let funding_out_idx: u32 = 0;
@@ -286,7 +288,8 @@ mod tests {
 
     #[test]
     fn test_verify_slashing_tx_adaptor_sig() {
-        let (btc_del, params) = get_btc_delegation_and_params();
+        let btc_del = get_btc_delegation();
+        let params = get_params();
 
         let staking_tx: Transaction = deserialize(&btc_del.staking_tx).unwrap();
         let slashing_tx: Transaction = deserialize(&btc_del.slashing_tx).unwrap();
@@ -334,7 +337,8 @@ mod tests {
 
     #[test]
     fn test_verify_unbonding_slashing_tx_adaptor_sig() {
-        let (btc_del, params) = get_btc_delegation_and_params();
+        let btc_del = get_btc_delegation();
+        let params = get_params();
 
         let btc_undel = btc_del.btc_undelegation.unwrap();
         let unbonding_tx: Transaction = deserialize(&btc_undel.unbonding_tx).unwrap();
diff --git a/packages/test-utils/src/lib.rs b/packages/test-utils/src/lib.rs
index ffe85017..0f434a7b 100644
--- a/packages/test-utils/src/lib.rs
+++ b/packages/test-utils/src/lib.rs
@@ -140,16 +140,16 @@ pub fn get_btc_timestamp_and_headers() -> (BtcTimestamp, HashMap<BlockHash, Bloc
     (btc_ts, header_map)
 }
 
-pub fn get_btc_delegation_and_params() -> (BtcDelegation, BtcStakingParams) {
-    let btc_del_path = find_testdata_path().join(BTC_DELEGATION_DATA);
-    let btc_del_data: &[u8] = &fs::read(btc_del_path).unwrap();
-    let btc_del = BtcDelegation::decode(btc_del_data).unwrap();
-
+pub fn get_params() -> BtcStakingParams {
     let params_path = find_testdata_path().join(PARAMS_DATA);
     let params_data: &[u8] = &fs::read(params_path).unwrap();
-    let params = BtcStakingParams::decode(params_data).unwrap();
+    BtcStakingParams::decode(params_data).unwrap()
+}
 
-    (btc_del, params)
+pub fn get_btc_delegation() -> BtcDelegation {
+    let btc_del_path = find_testdata_path().join(BTC_DELEGATION_DATA);
+    let btc_del_data: &[u8] = &fs::read(btc_del_path).unwrap();
+    BtcDelegation::decode(btc_del_data).unwrap()
 }
 
 pub fn get_pub_rand_commit() -> MsgCommitPubRandList {