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

iota_cluster_test/
lib.rs
