iota_cluster_test/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use cluster::{Cluster, ClusterFactory};
9use config::ClusterTestOpt;
10use futures::{StreamExt, stream::FuturesUnordered};
11use helper::ObjectChecker;
12use iota_faucet::CoinInfo;
13use iota_json_rpc_types::{
14    IotaExecutionStatus, IotaTransactionBlockEffectsAPI, IotaTransactionBlockResponse,
15    IotaTransactionBlockResponseOptions, TransactionBlockBytes,
16};
17use iota_sdk::{IotaClient, wallet_context::WalletContext};
18use iota_test_transaction_builder::batch_make_transfer_transactions;
19use iota_types::{
20    base_types::{IotaAddress, TransactionDigest},
21    gas_coin::GasCoin,
22    iota_system_state::iota_system_state_summary::IotaSystemStateSummary,
23    object::Owner,
24    quorum_driver_types::ExecuteTransactionRequestType,
25    transaction::{Transaction, TransactionData},
26};
27use jsonrpsee::{
28    core::{client::ClientT, params::ArrayParams},
29    http_client::HttpClientBuilder,
30};
31use test_case::{
32    coin_index_test::CoinIndexTest, coin_merge_split_test::CoinMergeSplitTest,
33    fullnode_build_publish_transaction_test::FullNodeBuildPublishTransactionTest,
34    fullnode_execute_transaction_test::FullNodeExecuteTransactionTest,
35    native_transfer_test::NativeTransferTest, random_beacon_test::RandomBeaconTest,
36    shared_object_test::SharedCounterTest,
37};
38use tokio::time::{self, Duration};
39use tracing::{error, info};
40use wallet_client::WalletClient;
41
42use crate::faucet::{FaucetClient, FaucetClientFactory};
43
44pub mod cluster;
45pub mod config;
46pub mod faucet;
47pub mod helper;
48pub mod test_case;
49pub mod wallet_client;
50
51pub use iota_genesis_builder::SnapshotUrl as MigrationSnapshotUrl;
52
53pub struct TestContext {
54    /// Cluster handle that allows access to various components in a cluster
55    cluster: Box<dyn Cluster + Sync + Send>,
56    /// Client that provides wallet context and fullnode access
57    client: WalletClient,
58    /// Facuet client that provides faucet access to a test
59    faucet: Arc<dyn FaucetClient + Sync + Send>,
60}
61
62impl TestContext {
63    async fn get_iota_from_faucet(&self, minimum_coins: Option<usize>) -> Vec<GasCoin> {
64        let addr = self.get_wallet_address();
65        let faucet_response = self.faucet.request_iota_coins(addr).await;
66
67        let coin_info = faucet_response
68            .transferred_gas_objects
69            .iter()
70            .map(|coin_info| coin_info.transfer_tx_digest)
71            .collect::<Vec<_>>();
72        self.let_fullnode_sync(coin_info, 5).await;
73
74        let gas_coins = self
75            .check_owner_and_into_gas_coin(faucet_response.transferred_gas_objects, addr)
76            .await;
77
78        let minimum_coins = minimum_coins.unwrap_or(1);
79
80        if gas_coins.len() < minimum_coins {
81            panic!(
82                "Expect to get at least {minimum_coins} IOTA Coins for address {addr}, but only got {}",
83                gas_coins.len()
84            )
85        }
86
87        gas_coins
88    }
89
90    fn get_context(&self) -> &WalletClient {
91        &self.client
92    }
93
94    fn get_fullnode_client(&self) -> &IotaClient {
95        self.client.get_fullnode_client()
96    }
97
98    fn clone_fullnode_client(&self) -> IotaClient {
99        self.client.get_fullnode_client().clone()
100    }
101
102    fn get_fullnode_rpc_url(&self) -> &str {
103        self.cluster.fullnode_url()
104    }
105
106    fn get_wallet(&self) -> &WalletContext {
107        self.client.get_wallet()
108    }
109
110    async fn get_latest_iota_system_state(&self) -> IotaSystemStateSummary {
111        self.client
112            .get_fullnode_client()
113            .governance_api()
114            .get_latest_iota_system_state()
115            .await
116            .unwrap()
117    }
118
119    async fn get_reference_gas_price(&self) -> u64 {
120        self.client
121            .get_fullnode_client()
122            .governance_api()
123            .get_reference_gas_price()
124            .await
125            .unwrap()
126    }
127
128    fn get_wallet_address(&self) -> IotaAddress {
129        self.client.get_wallet_address()
130    }
131
132    /// See `make_transactions_with_wallet_context` for potential caveats
133    /// of this helper function.
134    pub async fn make_transactions(&self, max_txn_num: usize) -> Vec<Transaction> {
135        batch_make_transfer_transactions(self.get_wallet(), max_txn_num).await
136    }
137
138    pub async fn build_transaction_remotely(
139        &self,
140        method: &str,
141        params: ArrayParams,
142    ) -> anyhow::Result<TransactionData> {
143        let fn_rpc_url = self.get_fullnode_rpc_url();
144        // TODO cache this?
145        let rpc_client = HttpClientBuilder::default().build(fn_rpc_url)?;
146
147        TransactionBlockBytes::to_data(rpc_client.request(method, params).await?)
148    }
149
150    async fn sign_and_execute(
151        &self,
152        txn_data: TransactionData,
153        desc: &str,
154    ) -> IotaTransactionBlockResponse {
155        let signature = self.get_context().sign(&txn_data, desc);
156        let resp = self
157            .get_fullnode_client()
158            .quorum_driver_api()
159            .execute_transaction_block(
160                Transaction::from_data(txn_data, vec![signature]),
161                IotaTransactionBlockResponseOptions::new()
162                    .with_object_changes()
163                    .with_balance_changes()
164                    .with_effects()
165                    .with_events(),
166                Some(ExecuteTransactionRequestType::WaitForLocalExecution),
167            )
168            .await
169            .unwrap_or_else(|e| panic!("Failed to execute transaction for {desc}. {e}"));
170        assert!(
171            matches!(
172                resp.effects.as_ref().unwrap().status(),
173                IotaExecutionStatus::Success
174            ),
175            "Failed to execute transaction for {desc}: {resp:?}"
176        );
177        resp
178    }
179
180    pub async fn setup(options: ClusterTestOpt) -> Result<Self, anyhow::Error> {
181        let cluster = ClusterFactory::start(&options).await?;
182        let wallet_client = WalletClient::new_from_cluster(&cluster).await;
183        let faucet = FaucetClientFactory::new_from_cluster(&cluster).await;
184        Ok(Self {
185            cluster,
186            client: wallet_client,
187            faucet,
188        })
189    }
190
191    // TODO: figure out a more efficient way to test a local cluster
192    // A potential way to do this is to subscribe to txns from fullnode
193    // when the feature is ready
194    pub async fn let_fullnode_sync(&self, digests: Vec<TransactionDigest>, timeout_sec: u64) {
195        let mut futures = FuturesUnordered::new();
196        for digest in digests.clone() {
197            let task = self.get_tx_with_retry_times(digest, 1);
198            futures.push(Box::pin(task));
199        }
200        let mut sleep = Box::pin(time::sleep(Duration::from_secs(timeout_sec)));
201
202        loop {
203            tokio::select! {
204                _ = &mut sleep => {
205                    panic!("Fullnode does not know all of {digests:?} after {timeout_sec} secs.");
206                }
207                res = futures.next() => {
208                    match res {
209                        Some((true, _, _)) => {},
210                        Some((false, digest, retry_times)) => {
211                            let task = self.get_tx_with_retry_times(digest, retry_times);
212                            futures.push(Box::pin(task));
213                        },
214                        None => break, // all txns appear on fullnode, mission completed
215                    }
216                }
217            }
218        }
219    }
220
221    async fn get_tx_with_retry_times(
222        &self,
223        digest: TransactionDigest,
224        retry_times: u64,
225    ) -> (bool, TransactionDigest, u64) {
226        match self
227            .client
228            .get_fullnode_client()
229            .read_api()
230            .get_transaction_with_options(digest, IotaTransactionBlockResponseOptions::new())
231            .await
232        {
233            Ok(_) => (true, digest, retry_times),
234            Err(_) => {
235                time::sleep(Duration::from_millis(300 * retry_times)).await;
236                (false, digest, retry_times + 1)
237            }
238        }
239    }
240
241    async fn check_owner_and_into_gas_coin(
242        &self,
243        coin_info: Vec<CoinInfo>,
244        owner: IotaAddress,
245    ) -> Vec<GasCoin> {
246        futures::future::join_all(
247            coin_info
248                .iter()
249                .map(|coin_info| {
250                    ObjectChecker::new(coin_info.id)
251                        .owner(Owner::AddressOwner(owner))
252                        .check_into_gas_coin(self.get_fullnode_client())
253                })
254                .collect::<Vec<_>>(),
255        )
256        .await
257        .into_iter()
258        .collect::<Vec<_>>()
259    }
260}
261
262pub struct TestCase<'a> {
263    test_case: Box<dyn TestCaseImpl + 'a>,
264}
265
266impl<'a> TestCase<'a> {
267    pub fn new(test_case: impl TestCaseImpl + 'a) -> Self {
268        TestCase {
269            test_case: (Box::new(test_case)),
270        }
271    }
272
273    pub async fn run(self, ctx: &mut TestContext) -> bool {
274        let test_name = self.test_case.name();
275        info!("Running test {}.", test_name);
276
277        // TODO: unwind panic and fail gracefully?
278
279        match self.test_case.run(ctx).await {
280            Ok(()) => {
281                info!("Test {test_name} succeeded.");
282                true
283            }
284            Err(e) => {
285                error!("Test {test_name} failed with error: {e}.");
286                false
287            }
288        }
289    }
290}
291
292#[async_trait]
293pub trait TestCaseImpl {
294    fn name(&self) -> &'static str;
295    fn description(&self) -> &'static str;
296    async fn run(&self, ctx: &mut TestContext) -> Result<(), anyhow::Error>;
297}
298
299pub struct ClusterTest;
300
301impl ClusterTest {
302    pub async fn run(options: ClusterTestOpt) {
303        let mut ctx = TestContext::setup(options)
304            .await
305            .unwrap_or_else(|e| panic!("Failed to set up TestContext, e: {e}"));
306
307        // TODO: collect tests from each test_case file instead.
308        let tests = vec![
309            TestCase::new(NativeTransferTest {}),
310            TestCase::new(CoinMergeSplitTest {}),
311            TestCase::new(SharedCounterTest {}),
312            TestCase::new(FullNodeExecuteTransactionTest {}),
313            TestCase::new(FullNodeBuildPublishTransactionTest {}),
314            TestCase::new(CoinIndexTest {}),
315            TestCase::new(RandomBeaconTest {}),
316        ];
317
318        // TODO: improve the runner parallelism for efficiency
319        // For now we run tests serially
320        let mut success_cnt = 0;
321        let total_cnt = tests.len() as i32;
322        for t in tests {
323            let is_success = t.run(&mut ctx).await as i32;
324            success_cnt += is_success;
325        }
326        if success_cnt < total_cnt {
327            // If any test failed, panic to bubble up the signal
328            panic!("{success_cnt} of {total_cnt} tests passed.");
329        }
330        info!("{success_cnt} of {total_cnt} tests passed.");
331    }
332}
iota_cluster_test/lib.rs

iota_cluster_test/
lib.rs
