test: add more test

nervosnetwork · May 8, 2024 · 4778660 · 4778660
1 parent c76e8be
commit 4778660
Show file tree

Hide file tree

Showing 13 changed files with 770 additions and 127 deletions.
diff --git a/examples/omnilock_examples.rs b/examples/omnilock_examples.rs
@@ -51,6 +51,7 @@ fn main() -> Result<(), Box<dyn StdErr>> {
         &SignContexts::new_omnilock(
             [secp256k1::SecretKey::from_slice(private_key.as_bytes())?].to_vec(),
             omni_cfg,
+            ckb_sdk::unlock::OmniUnlockMode::Normal,
         ),
     )?;
 

diff --git a/examples/omnilock_multisig_example.rs b/examples/omnilock_multisig_example.rs
@@ -62,6 +62,7 @@ fn main() -> Result<(), Box<dyn StdErr>> {
         &SignContexts::new_omnilock(
             [secp256k1::SecretKey::from_slice(private_key.as_bytes())?].to_vec(),
             omni_cfg.clone(),
+            ckb_sdk::unlock::OmniUnlockMode::Normal,
         ),
     )?;
     let private_key = h256!("0x4fd809631a6aa6e3bb378dd65eae5d71df895a82c91a615a1e8264741515c79c");
@@ -70,6 +71,7 @@ fn main() -> Result<(), Box<dyn StdErr>> {
         &SignContexts::new_omnilock(
             [secp256k1::SecretKey::from_slice(private_key.as_bytes())?].to_vec(),
             omni_cfg,
+            ckb_sdk::unlock::OmniUnlockMode::Normal,
         ),
     )?;
 

diff --git a/src/core/advanced_builders.rs b/src/core/advanced_builders.rs
@@ -232,7 +232,8 @@ impl TransactionBuilder {
         if witness_data.is_empty() {
             WitnessArgs::default()
         } else {
-            WitnessArgs::from_slice(witness_data.as_ref()).unwrap()
+            WitnessArgs::from_slice(witness_data.as_ref())
+                .expect("WitnessArgs expected but failed to decode ")
         }
     }
 

diff --git a/src/tests/transaction/omnilock.rs b/src/tests/transaction/omnilock.rs
diff --git a/src/tests/tx_builder/mod.rs b/src/tests/tx_builder/mod.rs
@@ -3,6 +3,6 @@ mod cheque;
 mod cycle;
 mod dao;
 mod omni_lock;
-mod omni_lock_util;
+pub mod omni_lock_util;
 mod transfer;
 mod udt;
diff --git a/src/transaction/builder/mod.rs b/src/transaction/builder/mod.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use super::{handler::HandlerContexts, input::TransactionInput};
 use crate::{
     core::TransactionBuilder,
-    traits::CellCollectorError,
+    traits::{CellCollectorError, LiveCell},
     transaction::TransactionBuilderConfiguration,
     tx_builder::{BalanceTxCapacityError, TxBuilderError},
     ScriptGroup, TransactionWithScriptGroups,
@@ -143,6 +143,7 @@ fn inner_build<
     input_iter: I,
     configuration: &TransactionBuilderConfiguration,
     contexts: &HandlerContexts,
+    rc_cells: Vec<LiveCell>,
 ) -> Result<TransactionWithScriptGroups, TxBuilderError> {
     let mut lock_groups: HashMap<Byte32, ScriptGroup> = HashMap::default();
     let mut type_groups: HashMap<Byte32, ScriptGroup> = HashMap::default();
@@ -159,10 +160,43 @@ fn inner_build<
         }
     }
 
+    let rc_len = rc_cells.len();
+
+    for (input_index, input) in rc_cells.into_iter().enumerate() {
+        let input = TransactionInput::new(input, 0);
+        tx.input(input.cell_input());
+        tx.witness(packed::Bytes::default());
+
+        inputs.insert(
+            input.live_cell.out_point.clone(),
+            (
+                input.live_cell.output.clone(),
+                input.live_cell.output_data.clone(),
+            ),
+        );
+
+        let previous_output = input.previous_output();
+        let lock_script = previous_output.lock();
+        lock_groups
+            .entry(lock_script.calc_script_hash())
+            .or_insert_with(|| ScriptGroup::from_lock_script(&lock_script))
+            .input_indices
+            .push(input_index);
+
+        if let Some(type_script) = previous_output.type_().to_opt() {
+            type_groups
+                .entry(type_script.calc_script_hash())
+                .or_insert_with(|| ScriptGroup::from_type_script(&type_script))
+                .input_indices
+                .push(input_index);
+        }
+    }
+
     // setup change output and data
     change_builder.init(&mut tx);
     // collect inputs
-    for (input_index, input) in input_iter.enumerate() {
+    for (mut input_index, input) in input_iter.enumerate() {
+        input_index = input_index + rc_len;
         let input = input?;
         tx.input(input.cell_input());
         tx.witness(packed::Bytes::default());

diff --git a/src/transaction/builder/simple.rs b/src/transaction/builder/simple.rs
@@ -1,5 +1,6 @@
 use crate::{
     core::TransactionBuilder,
+    traits::LiveCell,
     transaction::{
         handler::HandlerContexts, input::InputIterator, TransactionBuilderConfiguration,
     },
@@ -24,6 +25,7 @@ pub struct SimpleTransactionBuilder {
     input_iter: InputIterator,
     /// The inner transaction builder
     tx: TransactionBuilder,
+    rc_cells: Vec<LiveCell>,
 }
 
 impl SimpleTransactionBuilder {
@@ -37,9 +39,14 @@ impl SimpleTransactionBuilder {
             configuration,
             input_iter,
             tx: TransactionBuilder::default(),
+            rc_cells: Vec::new(),
         }
     }
 
+    pub fn set_rc_cells(&mut self, rc_cells: Vec<LiveCell>) {
+        self.rc_cells = rc_cells
+    }
+
     /// Update the change lock script.
     pub fn set_change_lock(&mut self, lock_script: Script) {
         self.change_lock = lock_script;
@@ -71,6 +78,7 @@ impl CkbTransactionBuilder for SimpleTransactionBuilder {
             configuration,
             input_iter,
             tx,
+            rc_cells,
         } = self;
 
         let change_builder = DefaultChangeBuilder {
@@ -79,6 +87,13 @@ impl CkbTransactionBuilder for SimpleTransactionBuilder {
             inputs: Vec::new(),
         };
 
-        inner_build(tx, change_builder, input_iter, &configuration, contexts)
+        inner_build(
+            tx,
+            change_builder,
+            input_iter,
+            &configuration,
+            contexts,
+            rc_cells,
+        )
     }
 }
diff --git a/src/transaction/builder/sudt.rs b/src/transaction/builder/sudt.rs
@@ -128,7 +128,14 @@ impl CkbTransactionBuilder for SudtTransactionBuilder {
         };
 
         if owner_mode {
-            inner_build(tx, change_builder, input_iter, &configuration, contexts)
+            inner_build(
+                tx,
+                change_builder,
+                input_iter,
+                &configuration,
+                contexts,
+                Default::default(),
+            )
         } else {
             let sudt_type_script =
                 build_sudt_type_script(configuration.network_info(), &sudt_owner_lock_script);
@@ -154,7 +161,14 @@ impl CkbTransactionBuilder for SudtTransactionBuilder {
                         .to_le_bytes()
                         .pack();
                     tx.set_output_data(tx.outputs_data.len() - 1, change_output_data);
-                    return inner_build(tx, change_builder, input_iter, &configuration, contexts);
+                    return inner_build(
+                        tx,
+                        change_builder,
+                        input_iter,
+                        &configuration,
+                        contexts,
+                        Default::default(),
+                    );
                 }
             }
 

diff --git a/src/transaction/handler/mod.rs b/src/transaction/handler/mod.rs
@@ -77,6 +77,7 @@ impl HandlerContexts {
     }
 }
 
+#[allow(unused_macros)]
 macro_rules! cell_dep {
     ($hash: literal, $idx: expr, $dep_type: expr) => {{
         let out_point = ckb_types::packed::OutPoint::new_builder()
@@ -90,4 +91,5 @@ macro_rules! cell_dep {
     }};
 }
 
+#[allow(unused_imports)]
 pub(crate) use cell_dep;
diff --git a/src/transaction/handler/omnilock.rs b/src/transaction/handler/omnilock.rs
@@ -1,11 +1,10 @@
 use ckb_types::{
-    core::DepType,
     h256,
     packed::{CellDep, OutPoint, Script},
     prelude::{Builder, Entity, Pack},
 };
 
-use super::{cell_dep, HandlerContext, ScriptHandler};
+use super::{HandlerContext, ScriptHandler};
 use crate::{
     core::TransactionBuilder,
     tx_builder::TxBuilderError,
@@ -20,8 +19,6 @@ use lazy_static::lazy_static;
 
 pub struct OmnilockScriptHandler {
     cell_deps: Vec<CellDep>,
-    sighash_dep: CellDep,
-    multisig_dep: CellDep,
     lock_script_id: ScriptId,
 }
 
@@ -41,6 +38,11 @@ impl OmnilockScriptContext {
             rpc_url,
         }
     }
+
+    pub fn unlock_mode(mut self, unlock_mode: OmniUnlockMode) -> Self {
+        self.unlock_mode = unlock_mode;
+        self
+    }
 }
 
 impl HandlerContext for OmnilockScriptContext {}
@@ -53,8 +55,6 @@ impl OmnilockScriptHandler {
     pub fn new_with_network(network: &NetworkInfo) -> Result<Self, TxBuilderError> {
         let mut ret = Self {
             cell_deps: vec![],
-            sighash_dep: Default::default(),
-            multisig_dep: Default::default(),
             lock_script_id: ScriptId::default(),
         };
         ret.init(network)?;
@@ -65,6 +65,10 @@ impl OmnilockScriptHandler {
         self.cell_deps = cell_deps;
     }
 
+    pub fn insert_cell_dep(&mut self, cell_dep: CellDep) {
+        self.cell_deps.push(cell_dep)
+    }
+
     pub fn set_lock_script_id(&mut self, lock_script_id: ScriptId) {
         self.lock_script_id = lock_script_id;
     }
@@ -128,40 +132,8 @@ impl ScriptHandler for OmnilockScriptHandler {
 
     fn init(&mut self, network: &NetworkInfo) -> Result<(), TxBuilderError> {
         if network.network_type == NetworkType::Mainnet {
-            self.sighash_dep = cell_dep!(
-                "0x71a7ba8fc96349fea0ed3a5c47992e3b4084b031a42264a018e0072e8172e46c",
-                0u32,
-                DepType::DepGroup
-            );
-            self.multisig_dep = cell_dep!(
-                "0x71a7ba8fc96349fea0ed3a5c47992e3b4084b031a42264a018e0072e8172e46c",
-                1u32,
-                DepType::DepGroup
-            );
-            self.cell_deps.push(self.sighash_dep.clone());
-            self.cell_deps.push(cell_dep!(
-                "0xc76edf469816aa22f416503c38d0b533d2a018e253e379f134c3985b3472c842",
-                0u32,
-                DepType::Code
-            ));
             self.lock_script_id = MAINNET_OMNILOCK_SCRIPT_ID.clone();
         } else if network.network_type == NetworkType::Testnet {
-            self.sighash_dep = cell_dep!(
-                "0xf8de3bb47d055cdf460d93a2a6e1b05f7432f9777c8c474abf4eec1d4aee5d37",
-                0u32,
-                DepType::DepGroup
-            );
-            self.multisig_dep = cell_dep!(
-                "0xf8de3bb47d055cdf460d93a2a6e1b05f7432f9777c8c474abf4eec1d4aee5d37",
-                1u32,
-                DepType::DepGroup
-            );
-            self.cell_deps.push(self.sighash_dep.clone());
-            self.cell_deps.push(cell_dep!(
-                "0x3d4296df1bd2cc2bd3f483f61ab7ebeac462a2f336f2b944168fe6ba5d81c014",
-                0u32,
-                DepType::Code
-            ));
             self.lock_script_id = get_testnet_omnilock_script_id().clone();
         } else {
             return Err(TxBuilderError::UnsupportedNetworkType(network.network_type));

diff --git a/src/transaction/signer/mod.rs b/src/transaction/signer/mod.rs
@@ -4,7 +4,7 @@ use std::collections::HashMap;
 use crate::{
     constants,
     traits::TransactionDependencyProvider,
-    unlock::{MultisigConfig, OmniLockConfig, UnlockError},
+    unlock::{MultisigConfig, OmniLockConfig, OmniUnlockMode, UnlockError},
     NetworkInfo, NetworkType, ScriptGroup, ScriptId, TransactionWithScriptGroups,
 };
 
@@ -72,8 +72,13 @@ impl SignContexts {
         Ok(Self::new_multisig(key, multisig_config))
     }
 
-    pub fn new_omnilock(keys: Vec<secp256k1::SecretKey>, omnilock_config: OmniLockConfig) -> Self {
-        let omnilock_context = omnilock::OmnilockSignerContext::new(keys, omnilock_config);
+    pub fn new_omnilock(
+        keys: Vec<secp256k1::SecretKey>,
+        omnilock_config: OmniLockConfig,
+        unlock_mode: OmniUnlockMode,
+    ) -> Self {
+        let omnilock_context =
+            omnilock::OmnilockSignerContext::new(keys, omnilock_config).unlock_mode(unlock_mode);
         Self {
             contexts: vec![Box::new(omnilock_context)],
         }
@@ -82,9 +87,11 @@ impl SignContexts {
     pub fn new_omnilock_solana(
         key: ed25519_dalek::SigningKey,
         omnilock_config: OmniLockConfig,
+        unlock_mode: OmniUnlockMode,
     ) -> Self {
         let omnilock_context =
-            omnilock::OmnilockSignerContext::new_with_ed25519_key(key, omnilock_config);
+            omnilock::OmnilockSignerContext::new_with_ed25519_key(key, omnilock_config)
+                .unlock_mode(unlock_mode);
         Self {
             contexts: vec![Box::new(omnilock_context)],
         }
@@ -93,9 +100,11 @@ impl SignContexts {
     pub fn new_omnilock_exec_dl_custom<T: crate::traits::Signer + 'static>(
         signer: T,
         omnilock_config: OmniLockConfig,
+        unlock_mode: OmniUnlockMode,
     ) -> Self {
         let omnilock_context =
-            omnilock::OmnilockSignerContext::new_with_dl_exec_signer(signer, omnilock_config);
+            omnilock::OmnilockSignerContext::new_with_dl_exec_signer(signer, omnilock_config)
+                .unlock_mode(unlock_mode);
         Self {
             contexts: vec![Box::new(omnilock_context)],
         }

diff --git a/src/transaction/signer/omnilock.rs b/src/transaction/signer/omnilock.rs
@@ -53,6 +53,12 @@ impl OmnilockSignerContext {
         }
     }
 
+    /// Default is Normal
+    pub fn unlock_mode(mut self, unlock_mode: OmniUnlockMode) -> Self {
+        self.unlock_mode = unlock_mode;
+        self
+    }
+
     pub fn build_omnilock_unlocker(&self) -> OmniLockUnlocker {
         let signer: Box<dyn Signer> = match self.cfg.id().flag() {
             IdentityFlag::Ethereum | IdentityFlag::EthereumDisplaying | IdentityFlag::Tron => {