test_cluster/

lib.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use fastcrypto_zkp::bn254::zk_login::JwkId;
use futures::future::join_all;
use jsonrpsee::http_client::{HttpClient, HttpClientBuilder};
use mysten_common::ZipDebugEqIteratorExt;
use mysten_common::fatal;
use rand::{Rng, distributions::*, rngs::OsRng, seq::SliceRandom};
use std::net::SocketAddr;
use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use sui_config::genesis::Genesis;
use sui_config::node::FundsWithdrawSchedulerType;
use sui_config::node::{
    AuthorityOverloadConfig, ConsensusTransactionPoolConfig, DBCheckpointConfig, RunWithRange,
};
use sui_config::{Config, ExecutionCacheConfig, SUI_CLIENT_CONFIG, SUI_NETWORK_CONFIG};
use sui_config::{NodeConfig, PersistedConfig, SUI_KEYSTORE_FILENAME};
use sui_core::authority_aggregator::AuthorityAggregator;
use sui_core::authority_client::NetworkAuthorityClient;
use sui_core::transaction_driver::SubmitTransactionOptions;
use sui_keys::keystore::{AccountKeystore, FileBasedKeystore, Keystore};
use sui_node::SuiNodeHandle;
use sui_protocol_config::{Chain, ProtocolVersion};
use sui_rpc_api::Client;
use sui_rpc_api::client::ExecutedTransaction;
use sui_sdk::sui_client_config::{SuiClientConfig, SuiEnv};
use sui_sdk::wallet_context::WalletContext;
use sui_sdk::{SuiClient, SuiClientBuilder};
use sui_swarm::memory::{Swarm, SwarmBuilder};
use sui_swarm_config::genesis_config::{
    AccountConfig, DEFAULT_GAS_AMOUNT, GenesisConfig, ValidatorGenesisConfig,
};
use sui_swarm_config::network_config::NetworkConfig;
use sui_swarm_config::network_config_builder::{
    FundsWithdrawSchedulerTypeConfig, GlobalStateHashV2EnabledCallback,
    GlobalStateHashV2EnabledConfig, ProtocolVersionsConfig, SupportedProtocolVersionsCallback,
    ValidatorObserverConfigCallback,
};
use sui_swarm_config::node_config_builder::{FullnodeConfigBuilder, ValidatorConfigBuilder};
use sui_test_transaction_builder::TestTransactionBuilder;
use sui_types::authenticator_state::get_authenticator_state;
use sui_types::base_types::ConciseableName;
use sui_types::base_types::{AuthorityName, ObjectID, ObjectRef, SuiAddress};
use sui_types::committee::CommitteeTrait;
use sui_types::committee::{Committee, EpochId};
use sui_types::crypto::KeypairTraits;
use sui_types::crypto::SuiKeyPair;
use sui_types::digests::{ChainIdentifier, TransactionDigest};
use sui_types::effects::TransactionEffectsAPI;
use sui_types::effects::{TransactionEffects, TransactionEvents};
use sui_types::error::{SuiErrorKind, SuiResult};
use sui_types::messages_grpc::{
    RawSubmitTxRequest, SubmitTxRequest, SubmitTxResult, SubmitTxType, WaitForEffectsRequest,
    WaitForEffectsResponse,
};
use sui_types::object::Object;
use sui_types::sui_system_state::SuiSystemState;
use sui_types::sui_system_state::SuiSystemStateTrait;
use sui_types::sui_system_state::epoch_start_sui_system_state::EpochStartSystemStateTrait;
use sui_types::supported_protocol_versions::SupportedProtocolVersions;
use sui_types::traffic_control::{PolicyConfig, RemoteFirewallConfig};
use sui_types::transaction::{Transaction, TransactionData};
use tokio::sync::broadcast;
use tokio::time::{Instant, timeout};
use tokio::{task::JoinHandle, time::sleep};
use tonic::IntoRequest;
use tracing::{error, info};

pub mod addr_balance_test_env;

const NUM_VALIDATOR: usize = 4;
// Keep direct test submissions bounded like the production TransactionOrchestrator.
const TRANSACTION_FINALITY_TIMEOUT: Duration = Duration::from_secs(90);

pub struct FullNodeHandle {
    pub sui_node: SuiNodeHandle,
    #[deprecated = "use grpc_client"]
    pub sui_client: SuiClient,
    #[deprecated = "use grpc_client"]
    pub rpc_client: HttpClient,
    pub grpc_client: Client,
    pub rpc_url: String,
}

impl FullNodeHandle {
    pub async fn new(sui_node: SuiNodeHandle, json_rpc_address: SocketAddr) -> Self {
        let rpc_url = format!("http://{}", json_rpc_address);
        let rpc_client = HttpClientBuilder::default().build(&rpc_url).unwrap();

        let sui_client = SuiClientBuilder::default().build(&rpc_url).await.unwrap();
        let grpc_client = Client::new(&rpc_url).unwrap();

        Self {
            sui_node,
            #[allow(deprecated)]
            sui_client,
            #[allow(deprecated)]
            rpc_client,
            grpc_client,
            rpc_url,
        }
    }
}

pub struct TestCluster {
    pub swarm: Swarm,
    pub wallet: WalletContext,
    pub fullnode_handle: FullNodeHandle,
}

impl TestCluster {
    #[deprecated = "use grpc_client()"]
    pub fn rpc_client(&self) -> &HttpClient {
        #[allow(deprecated)]
        &self.fullnode_handle.rpc_client
    }

    #[deprecated = "use grpc_client()"]
    pub fn sui_client(&self) -> &SuiClient {
        #[allow(deprecated)]
        &self.fullnode_handle.sui_client
    }

    pub fn grpc_client(&self) -> Client {
        self.fullnode_handle.grpc_client.clone()
    }

    pub fn rpc_url(&self) -> &str {
        &self.fullnode_handle.rpc_url
    }

    pub fn wallet(&mut self) -> &WalletContext {
        &self.wallet
    }

    pub fn wallet_mut(&mut self) -> &mut WalletContext {
        &mut self.wallet
    }

    pub fn get_addresses(&self) -> Vec<SuiAddress> {
        self.wallet.get_addresses()
    }

    // Helper function to get the 0th address in WalletContext
    pub fn get_address_0(&self) -> SuiAddress {
        self.get_addresses()[0]
    }

    // Helper function to get the 1st address in WalletContext
    pub fn get_address_1(&self) -> SuiAddress {
        self.get_addresses()[1]
    }

    // Helper function to get the 2nd address in WalletContext
    pub fn get_address_2(&self) -> SuiAddress {
        self.get_addresses()[2]
    }

    pub fn fullnode_config_builder(&self) -> FullnodeConfigBuilder {
        self.swarm.get_fullnode_config_builder()
    }

    pub fn committee(&self) -> Arc<Committee> {
        self.fullnode_handle
            .sui_node
            .with(|node| node.state().epoch_store_for_testing().committee().clone())
    }

    pub fn get_sui_system_state(&self) -> SuiSystemState {
        self.fullnode_handle.sui_node.with(|node| {
            node.state()
                .get_sui_system_state_object_for_testing()
                .unwrap()
        })
    }

    /// Convenience method to start a new fullnode in the test cluster.
    pub async fn spawn_new_fullnode(&mut self) -> FullNodeHandle {
        self.start_fullnode_from_config(
            self.fullnode_config_builder()
                .build(&mut OsRng, self.swarm.config()),
        )
        .await
    }

    /// The observer fullnode of the cluster, when built with `with_observer_fullnode()`.
    /// The cluster deliberately does not hold a node handle for the observer: a handle
    /// keeps the running instance (and its open stores) alive, which prevents the
    /// simulator from restarting the node after a crash. Acquire a fresh handle via
    /// `Node::get_node_handle()` at the point of use and drop it promptly.
    pub fn observer_node(&self) -> Option<&sui_swarm::memory::Node> {
        self.swarm.observer_nodes().next()
    }

    pub async fn start_fullnode_from_config(&mut self, config: NodeConfig) -> FullNodeHandle {
        let json_rpc_address = config.json_rpc_address;
        let node = self.swarm.spawn_new_node(config).await;
        FullNodeHandle::new(node, json_rpc_address).await
    }

    pub fn all_node_handles(&self) -> Vec<SuiNodeHandle> {
        self.swarm
            .all_nodes()
            .flat_map(|n| n.get_node_handle())
            .collect()
    }

    pub fn all_validator_handles(&self) -> Vec<SuiNodeHandle> {
        self.swarm
            .validator_nodes()
            .map(|n| n.get_node_handle().unwrap())
            .collect()
    }

    pub fn get_validator_pubkeys(&self) -> Vec<AuthorityName> {
        self.swarm.active_validators().map(|v| v.name()).collect()
    }

    pub fn get_genesis(&self) -> Genesis {
        self.swarm.config().genesis.clone()
    }

    pub fn stop_node(&self, name: &AuthorityName) {
        self.swarm.node(name).unwrap().stop();
    }

    pub async fn stop_all_validators(&self) {
        info!("Stopping all validators in the cluster");
        self.swarm.active_validators().for_each(|v| v.stop());
        tokio::time::sleep(Duration::from_secs(3)).await;
    }

    pub async fn start_all_validators(&self) {
        info!("Starting all validators in the cluster");
        for v in self.swarm.validator_nodes() {
            if v.is_running() {
                continue;
            }
            v.start().await.unwrap();
        }
        tokio::time::sleep(Duration::from_secs(3)).await;
    }

    pub async fn start_node(&self, name: &AuthorityName) {
        let node = self.swarm.node(name).unwrap();
        if node.is_running() {
            return;
        }
        node.start().await.unwrap();
    }

    pub async fn spawn_new_validator(
        &mut self,
        genesis_config: ValidatorGenesisConfig,
    ) -> SuiNodeHandle {
        let node_config = ValidatorConfigBuilder::new()
            .build(genesis_config, self.swarm.config().genesis.clone());
        self.swarm.spawn_new_node(node_config).await
    }

    pub fn random_node_restarter(self: &Arc<Self>) -> RandomNodeRestarter {
        RandomNodeRestarter::new(self.clone())
    }

    pub async fn get_reference_gas_price(&self) -> u64 {
        self.grpc_client()
            .get_reference_gas_price()
            .await
            .expect("failed to get reference gas price")
    }

    pub fn get_chain_identifier(&self) -> ChainIdentifier {
        ChainIdentifier::from(*self.swarm.config().genesis.checkpoint().digest())
    }

    pub async fn get_object_from_fullnode_store(&self, object_id: &ObjectID) -> Option<Object> {
        self.fullnode_handle
            .sui_node
            .with_async(|node| async { node.state().get_object(object_id) })
            .await
    }

    pub async fn get_latest_object_ref(&self, object_id: &ObjectID) -> ObjectRef {
        self.get_object_from_fullnode_store(object_id)
            .await
            .unwrap()
            .compute_object_reference()
    }

    pub async fn get_object_or_tombstone_from_fullnode_store(
        &self,
        object_id: ObjectID,
    ) -> ObjectRef {
        self.fullnode_handle
            .sui_node
            .state()
            .get_object_cache_reader()
            .get_latest_object_ref_or_tombstone(object_id)
            .unwrap()
    }

    pub async fn wait_for_run_with_range_shutdown_signal(&self) -> Option<RunWithRange> {
        self.wait_for_run_with_range_shutdown_signal_with_timeout(Duration::from_secs(60))
            .await
    }

    pub async fn wait_for_run_with_range_shutdown_signal_with_timeout(
        &self,
        timeout_dur: Duration,
    ) -> Option<RunWithRange> {
        let mut shutdown_channel_rx = self
            .fullnode_handle
            .sui_node
            .with(|node| node.subscribe_to_shutdown_channel());

        timeout(timeout_dur, async move {
            tokio::select! {
                msg = shutdown_channel_rx.recv() =>
                {
                    match msg {
                        Ok(Some(run_with_range)) => Some(run_with_range),
                        Ok(None) => None,
                        Err(e) => {
                            error!("failed recv from sui-node shutdown channel: {}", e);
                            None
                        },
                    }
                },
            }
        })
        .await
        .expect("Timed out waiting for cluster to hit target epoch and recv shutdown signal from sui-node")
    }

    pub async fn wait_for_protocol_version(
        &self,
        target_protocol_version: ProtocolVersion,
    ) -> SuiSystemState {
        self.wait_for_protocol_version_with_timeout(
            target_protocol_version,
            Duration::from_secs(60),
        )
        .await
    }

    pub async fn wait_for_protocol_version_with_timeout(
        &self,
        target_protocol_version: ProtocolVersion,
        timeout_dur: Duration,
    ) -> SuiSystemState {
        timeout(timeout_dur, async move {
            loop {
                let system_state = self.wait_for_epoch(None).await;
                if system_state.protocol_version() >= target_protocol_version.as_u64() {
                    return system_state;
                }
            }
        })
        .await
        .expect("Timed out waiting for cluster to target protocol version")
    }

    /// Ask 2f+1 validators to close epoch actively, and wait for the entire network to reach the next
    /// epoch. This requires waiting for both the fullnode and all validators to reach the next epoch.
    pub async fn trigger_reconfiguration(&self) {
        info!("Starting reconfiguration");
        let start = Instant::now();

        // Close epoch on 2f+1 validators.
        let cur_committee = self
            .fullnode_handle
            .sui_node
            .with(|node| node.state().clone_committee_for_testing());
        let mut cur_stake = 0;
        for node in self.swarm.active_validators() {
            node.get_node_handle()
                .unwrap()
                .with_async(|node| async {
                    node.close_epoch_for_testing().await.unwrap_or_else(|_| {
                        fatal!(
                            "Failed to close epoch for validator {:?}",
                            node.state().name
                        );
                    });
                    cur_stake += cur_committee.weight(&node.state().name);
                })
                .await;
            if cur_stake >= cur_committee.quorum_threshold() {
                break;
            }
        }
        info!("close_epoch complete after {:?}", start.elapsed());

        self.wait_for_epoch(Some(cur_committee.epoch + 1)).await;
        self.wait_for_epoch_all_nodes(cur_committee.epoch + 1).await;

        info!("reconfiguration complete after {:?}", start.elapsed());
    }

    /// To detect whether the network has reached such state, we use the fullnode as the
    /// source of truth, since a fullnode only does epoch transition when the network has
    /// done so.
    /// If target_epoch is specified, wait until the cluster reaches that epoch.
    /// If target_epoch is None, wait until the cluster reaches the next epoch.
    /// Note that this function does not guarantee that every node is at the target epoch.
    pub async fn wait_for_epoch(&self, target_epoch: Option<EpochId>) -> SuiSystemState {
        self.wait_for_epoch_with_timeout(target_epoch, Duration::from_secs(60))
            .await
    }

    pub async fn wait_for_epoch_on_node(
        &self,
        handle: &SuiNodeHandle,
        target_epoch: Option<EpochId>,
        timeout_dur: Duration,
    ) -> SuiSystemState {
        let mut epoch_rx = handle.with(|node| node.subscribe_to_epoch_change());

        let mut state = None;
        timeout(timeout_dur, async {
            let epoch = handle.with(|node| node.state().epoch_store_for_testing().epoch());
            if Some(epoch) == target_epoch {
                return handle.with(|node| node.state().get_sui_system_state_object_for_testing().unwrap());
            }
            while let Ok(system_state) = epoch_rx.recv().await {
                info!("received epoch {}", system_state.epoch());
                state = Some(system_state.clone());
                match target_epoch {
                    Some(target_epoch) if system_state.epoch() >= target_epoch => {
                        return system_state;
                    }
                    None => {
                        return system_state;
                    }
                    _ => (),
                }
            }
            unreachable!("Broken reconfig channel");
        })
        .await
        .unwrap_or_else(|_| {
            error!("Timed out waiting for cluster to reach epoch {target_epoch:?}");
            if let Some(state) = state {
                panic!("Timed out waiting for cluster to reach epoch {target_epoch:?}. Current epoch: {}", state.epoch());
            }
            panic!("Timed out waiting for cluster to target epoch {target_epoch:?}")
        })
    }

    pub async fn wait_for_epoch_with_timeout(
        &self,
        target_epoch: Option<EpochId>,
        timeout_dur: Duration,
    ) -> SuiSystemState {
        self.wait_for_epoch_on_node(&self.fullnode_handle.sui_node, target_epoch, timeout_dur)
            .await
    }

    pub async fn wait_for_epoch_all_nodes(&self, target_epoch: EpochId) {
        let handles: Vec<_> = self
            .swarm
            .all_nodes()
            .map(|node| node.get_node_handle().unwrap())
            .collect();
        let tasks: Vec<_> = handles
            .iter()
            .map(|handle| {
                handle.with_async(|node| async {
                    let mut retries = 0;
                    loop {
                        let epoch = node.state().epoch_store_for_testing().epoch();
                        if epoch == target_epoch {
                            if let Some(agg) = node.clone_authority_aggregator() {
                                // This is a fullnode, we need to wait for its auth aggregator to reconfigure as well.
                                if agg.committee.epoch() == target_epoch {
                                    break;
                                }
                            } else {
                                // This is a validator, we don't need to check the auth aggregator.
                                break;
                            }
                        }
                        tokio::time::sleep(Duration::from_secs(1)).await;
                        retries += 1;
                        if retries % 5 == 0 {
                            tracing::warn!(validator=?node.state().name.concise(), "Waiting for {:?} seconds to reach epoch {:?}. Currently at epoch {:?}", retries, target_epoch, epoch);
                        }
                    }
                })
            })
            .collect();

        timeout(Duration::from_secs(40), join_all(tasks))
            .await
            .expect("timed out waiting for reconfiguration to complete");
    }

    /// Waits until every node advances to a strictly higher epoch than the one it is at when this is
    /// called. Unlike `wait_for_epoch_all_nodes`, which matches a target epoch exactly, this only
    /// requires forward progress, so it is safe when the cluster catches up through several epochs at
    /// once (e.g. recovering from a stall) and would blow past any fixed target.
    pub async fn wait_for_next_epoch_all_nodes(&self) {
        let handles: Vec<_> = self
            .swarm
            .all_nodes()
            .map(|node| node.get_node_handle().unwrap())
            .collect();
        let tasks: Vec<_> = handles
            .iter()
            .map(|handle| {
                handle.with_async(|node| async {
                    let start_epoch = node.state().epoch_store_for_testing().epoch();
                    let mut retries = 0;
                    loop {
                        let epoch = node.state().epoch_store_for_testing().epoch();
                        if epoch > start_epoch {
                            if let Some(agg) = node.clone_authority_aggregator() {
                                // Fullnode: also wait for its auth aggregator to reconfigure.
                                if agg.committee.epoch() > start_epoch {
                                    break;
                                }
                            } else {
                                break;
                            }
                        }
                        tokio::time::sleep(Duration::from_secs(1)).await;
                        retries += 1;
                        if retries % 5 == 0 {
                            tracing::warn!(validator=?node.state().name.concise(), "Waiting {:?}s for an epoch beyond {:?}; currently at {:?}", retries, start_epoch, epoch);
                        }
                    }
                })
            })
            .collect();

        timeout(Duration::from_secs(40), join_all(tasks))
            .await
            .expect("timed out waiting for all nodes to advance an epoch");
    }

    pub fn subscribe_to_epoch_change(&self) -> broadcast::Receiver<SuiSystemState> {
        // fullnode_handle is not part of swarm and cannot be dropped / killed
        self.fullnode_handle
            .sui_node
            .with(|node| node.subscribe_to_epoch_change())
    }

    /// Upgrade the network protocol version, by restarting every validator with a new
    /// supported versions.
    /// Note that we don't restart the fullnode here, and it is assumed that the fulnode supports
    /// the entire version range.
    pub async fn update_validator_supported_versions(
        &self,
        new_supported_versions: SupportedProtocolVersions,
    ) {
        for authority in self.get_validator_pubkeys() {
            self.stop_node(&authority);
            tokio::time::sleep(Duration::from_millis(1000)).await;
            self.swarm
                .node(&authority)
                .unwrap()
                .config()
                .supported_protocol_versions = Some(new_supported_versions);
            self.start_node(&authority).await;
            info!("Restarted validator {}", authority);
        }
    }

    /// Wait for all nodes in the network to upgrade to `protocol_version`.
    pub async fn wait_for_all_nodes_upgrade_to(&self, protocol_version: u64) {
        for h in self.all_node_handles() {
            h.with_async(|node| async {
                while node
                    .state()
                    .epoch_store_for_testing()
                    .epoch_start_state()
                    .protocol_version()
                    .as_u64()
                    != protocol_version
                {
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
            })
            .await;
        }
    }

    /// Wait until the on-chain authenticator state contains any active JWK.
    pub async fn wait_for_authenticator_state_update(&self) {
        self.wait_for_authenticator_state_update_for_providers(&[])
            .await;
    }

    /// Wait until the on-chain authenticator state contains all the given JWK ids.
    pub async fn wait_for_authenticator_state_update_for_providers(&self, jwk_ids: &[JwkId]) {
        timeout(Duration::from_secs(60), async {
            loop {
                let active: Vec<JwkId> = self.fullnode_handle.sui_node.with(|node| {
                    get_authenticator_state(node.state().get_object_store())
                        .ok()
                        .flatten()
                        .map(|state| {
                            state
                                .active_jwks
                                .iter()
                                .map(|active| active.jwk_id.clone())
                                .collect()
                        })
                        .unwrap_or_default()
                });
                let ready = if jwk_ids.is_empty() {
                    !active.is_empty()
                } else {
                    jwk_ids.iter().all(|id| active.contains(id))
                };
                if ready {
                    return;
                }
                tokio::time::sleep(Duration::from_millis(500)).await;
            }
        })
        .await
        .expect("Timed out waiting for authenticator state update");
    }

    /// Return the highest observed protocol version in the test cluster.
    pub fn highest_protocol_version(&self) -> ProtocolVersion {
        self.all_node_handles()
            .into_iter()
            .map(|h| {
                h.with(|node| {
                    node.state()
                        .epoch_store_for_testing()
                        .epoch_start_state()
                        .protocol_version()
                })
            })
            .max()
            .expect("at least one node must be up to get highest protocol version")
    }

    pub async fn test_transaction_builder(&self) -> TestTransactionBuilder {
        let (sender, gas) = self.wallet.get_one_gas_object().await.unwrap().unwrap();
        self.test_transaction_builder_with_gas_object(sender, gas)
            .await
    }

    pub async fn test_transaction_builder_with_sender(
        &self,
        sender: SuiAddress,
    ) -> TestTransactionBuilder {
        let gas = self
            .wallet
            .get_one_gas_object_owned_by_address(sender)
            .await
            .unwrap()
            .unwrap();
        self.test_transaction_builder_with_gas_object(sender, gas)
            .await
    }

    pub async fn test_transaction_builder_with_gas_object(
        &self,
        sender: SuiAddress,
        gas: ObjectRef,
    ) -> TestTransactionBuilder {
        let rgp = self.get_reference_gas_price().await;
        TestTransactionBuilder::new(sender, gas, rgp)
    }

    pub async fn sign_transaction(&self, tx_data: &TransactionData) -> Transaction {
        self.wallet.sign_transaction(tx_data).await
    }

    pub async fn sign_and_execute_transaction(
        &self,
        tx_data: &TransactionData,
    ) -> ExecutedTransaction {
        let tx = self.wallet.sign_transaction(tx_data).await;
        self.execute_transaction(tx).await
    }

    /// Sign and execute the transaction via direct validator submission, bypassing the fullnode.
    pub async fn sign_and_execute_transaction_directly(
        &self,
        tx_data: &TransactionData,
    ) -> SuiResult<(TransactionDigest, TransactionEffects)> {
        let mut res = self
            .sign_and_execute_txns_in_soft_bundle(std::slice::from_ref(tx_data))
            .await?;
        assert_eq!(res.len(), 1);
        Ok(res.pop().unwrap())
    }

    /// Execute an already-signed transaction via direct validator submission, bypassing the fullnode.
    pub async fn execute_transaction_directly(
        &self,
        tx: &Transaction,
    ) -> SuiResult<(TransactionDigest, TransactionEffects)> {
        let mut res = self
            .execute_signed_txns_in_soft_bundle(std::slice::from_ref(tx))
            .await?;
        assert_eq!(res.len(), 1);
        Ok(res.pop().unwrap())
    }

    /// Sign and execute multiple transactions in a soft bundle.
    /// Soft bundles allow submitting multiple transactions together with best-effort
    /// ordering if they use the same gas price. Transactions in a soft bundle can be
    /// individually rejected or deferred without affecting other transactions.
    ///
    /// NOTE: This is a simplified implementation that processes transactions individually.
    /// For true soft bundle submission, the test file should use the raw gRPC client directly
    /// with tonic, as shown in test_soft_bundle_different_gas_payers.
    pub async fn sign_and_execute_txns_in_soft_bundle(
        &self,
        txns: &[TransactionData],
    ) -> SuiResult<Vec<(TransactionDigest, TransactionEffects)>> {
        // Sign all transactions
        let signed_txs: Vec<Transaction> =
            futures::future::join_all(txns.iter().map(|tx| self.wallet.sign_transaction(tx))).await;

        self.execute_signed_txns_in_soft_bundle(&signed_txs).await
    }

    pub async fn execute_signed_txns_in_soft_bundle(
        &self,
        signed_txs: &[Transaction],
    ) -> SuiResult<Vec<(TransactionDigest, TransactionEffects)>> {
        let digests: Vec<_> = signed_txs.iter().map(|tx| *tx.digest()).collect();

        let request = RawSubmitTxRequest {
            transactions: signed_txs
                .iter()
                .map(|tx| bcs::to_bytes(tx).unwrap().into())
                .collect(),
            submit_type: SubmitTxType::SoftBundle.into(),
        };

        let agg = self.authority_aggregator();
        let clients = &agg.authority_clients;
        // Use seeded RNG for deterministic but varying validator selection in simtests
        let index = rand::thread_rng().gen_range(0..clients.len());
        let (_, safe_client) = clients.iter().nth(index).unwrap();
        let mut validator_client = safe_client
            .authority_client()
            .get_client_for_testing()
            .unwrap();

        let result = validator_client
            .submit_transaction(request.into_request())
            .await
            .map(tonic::Response::into_inner)?;
        assert_eq!(result.results.len(), signed_txs.len());

        let mut executed_results = vec![None; signed_txs.len()];
        let mut submitted_positions = Vec::new();
        for (index, raw_result) in result.results.into_iter().enumerate() {
            let submit_result: SubmitTxResult = raw_result.try_into()?;
            match submit_result {
                SubmitTxResult::Executed { details, .. } => {
                    let data = details.ok_or_else(|| SuiErrorKind::GenericAuthorityError {
                        error: "Expected execution details".to_string(),
                    })?;
                    executed_results[index] = Some((digests[index], data.effects));
                }
                SubmitTxResult::Rejected { error } => {
                    return Err(error);
                }
                SubmitTxResult::Submitted { consensus_position } => {
                    submitted_positions.push((index, consensus_position));
                }
            }
        }

        let wait_futures: Vec<_> = submitted_positions
            .iter()
            .map(|(index, position)| {
                let request = WaitForEffectsRequest {
                    transaction_digest: Some(digests[*index]),
                    consensus_position: Some(*position),
                    include_details: true,
                    ping_type: None,
                };
                safe_client.wait_for_effects(request, None)
            })
            .collect();

        let wait_responses = join_all(wait_futures).await;
        for ((index, _), response) in submitted_positions.into_iter().zip_debug_eq(wait_responses) {
            match response? {
                WaitForEffectsResponse::Executed { details, .. } => {
                    let data = details.ok_or_else(|| SuiErrorKind::GenericAuthorityError {
                        error: "Expected execution details".to_string(),
                    })?;
                    executed_results[index] = Some((digests[index], data.effects));
                }
                WaitForEffectsResponse::Rejected { error } => {
                    return Err(error.unwrap_or_else(|| {
                        SuiErrorKind::GenericAuthorityError {
                            error: "Transaction was rejected".to_string(),
                        }
                        .into()
                    }));
                }
                WaitForEffectsResponse::Expired { .. } => {
                    return Err(SuiErrorKind::TransactionExpired.into());
                }
            }
        }

        // Effects were already obtained from the validator above. Wait for the
        // rpc fullnode to settle the transactions in an executed checkpoint and
        // for its embedded rpc-store index to catch up, so callers that query
        // the fullnode (e.g. RPC for owned objects or balances) after this
        // returns observe these transactions. The embedded indexer follows the
        // tip asynchronously and is not a blocker for execution, so the index
        // wait is required for read-after-write consistency.
        self.wait_for_tx_settlement(&digests).await;

        executed_results
            .into_iter()
            .map(|result| {
                result.ok_or_else(|| {
                    SuiErrorKind::GenericAuthorityError {
                        error: "Missing execution result".to_string(),
                    }
                    .into()
                })
            })
            .collect()
    }

    /// Execute signed transactions in a soft bundle and return results for each transaction.
    /// Unlike `execute_signed_txns_in_soft_bundle`, this method handles conflicting transactions
    /// where some may be executed and others rejected.
    ///
    /// Returns a vector of (digest, WaitForEffectsResponse) for each transaction.
    pub async fn execute_soft_bundle_with_conflicts(
        &self,
        signed_txs: &[Transaction],
    ) -> SuiResult<Vec<(TransactionDigest, WaitForEffectsResponse)>> {
        let digests: Vec<_> = signed_txs.iter().map(|tx| *tx.digest()).collect();

        let request = RawSubmitTxRequest {
            transactions: signed_txs
                .iter()
                .map(|tx| bcs::to_bytes(tx).unwrap().into())
                .collect(),
            submit_type: SubmitTxType::SoftBundle.into(),
        };

        let authority_aggregator = self.authority_aggregator();
        let (_, safe_client) = authority_aggregator
            .authority_clients
            .iter()
            .next()
            .unwrap();
        let mut validator_client = safe_client
            .authority_client()
            .get_client_for_testing()
            .unwrap();

        let result = validator_client
            .submit_transaction(request.into_request())
            .await
            .map(tonic::Response::into_inner)?;
        assert_eq!(result.results.len(), signed_txs.len());

        // Extract consensus positions from submission results
        let mut consensus_positions = Vec::new();
        for (i, raw_result) in result.results.iter().enumerate() {
            let submit_result: SubmitTxResult = raw_result.clone().try_into()?;
            match submit_result {
                SubmitTxResult::Submitted { consensus_position } => {
                    consensus_positions.push(consensus_position);
                }
                SubmitTxResult::Executed { .. } => {
                    panic!(
                        "Transaction {} was already executed during submission",
                        i + 1
                    );
                }
                SubmitTxResult::Rejected { error } => {
                    return Err(error);
                }
            }
        }

        // Wait for effects using consensus positions
        let wait_futures: Vec<_> = digests
            .iter()
            .zip_debug_eq(consensus_positions.iter())
            .map(|(digest, position)| {
                let request = WaitForEffectsRequest {
                    transaction_digest: Some(*digest),
                    consensus_position: Some(*position),
                    include_details: false,
                    ping_type: None,
                };
                safe_client.wait_for_effects(request, None)
            })
            .collect();

        let responses = futures::future::join_all(wait_futures).await;

        let results: SuiResult<Vec<_>> = digests
            .into_iter()
            .zip_debug_eq(responses)
            .map(|(digest, response)| Ok((digest, response?)))
            .collect();

        results
    }

    pub async fn wait_for_tx_settlement(&self, digests: &[TransactionDigest]) {
        Self::wait_for_tx_settlement_on_handles(
            std::slice::from_ref(&self.fullnode_handle.sui_node),
            digests,
        )
        .await;
    }

    /// Like `wait_for_tx_settlement`, but waits on every node (all validators and fullnodes)
    /// instead of only the rpc fullnode. Nodes can execute the same checkpoint at different
    /// times, so the rpc fullnode having settled a transaction does not imply every validator
    /// has. Use this when a subsequent step may interact with an arbitrary validator (e.g. a
    /// soft bundle is submitted to a randomly chosen validator) and that validator must have
    /// already applied the settlement of `digests`.
    pub async fn wait_for_tx_settlement_all_nodes(&self, digests: &[TransactionDigest]) {
        let handles = self.all_node_handles();
        Self::wait_for_tx_settlement_on_handles(&handles, digests).await;
    }

    /// Waits, concurrently across `handles`, until every transaction in `digests` has settled
    /// on that node: the transaction's checkpoint is present in the node's store and that
    /// checkpoint has been executed (which is when the transaction's settlement is visible).
    async fn wait_for_tx_settlement_on_handles(
        handles: &[SuiNodeHandle],
        digests: &[TransactionDigest],
    ) {
        let waits = handles.iter().map(|handle| async move {
            let max_checkpoint_seq = handle
                .with_async(|node| async move {
                    let state = node.state();
                    // wait until the transactions are in checkpoints on this node
                    let checkpoint_seqs = state
                        .epoch_store_for_testing()
                        .transactions_executed_in_checkpoint_notify(digests.to_vec())
                        .await;

                    // then wait until the highest of those checkpoints is executed on this node
                    let max_checkpoint_seq = checkpoint_seqs.into_iter().max().unwrap();
                    state
                        .checkpoint_store
                        .notify_read_executed_checkpoint(max_checkpoint_seq)
                        .await;
                    max_checkpoint_seq
                })
                .await;

            // The embedded rpc-store indexes asynchronously, decoupled from
            // checkpoint execution, so a settled transaction is not yet visible
            // through the live index surface (owned objects, balances). Wait
            // for the live cohort to catch up so subsequent index reads
            // observe it.
            Self::wait_for_rpc_index_on_handle(handle, max_checkpoint_seq, false).await;
        });
        join_all(waits).await;
    }

    /// Wait until the embedded rpc-store on `handle` has indexed through
    /// `checkpoint`. No-op for a node without an embedded store (a validator,
    /// or a fullnode with indexing disabled).
    ///
    /// Unlike the legacy synchronous `rpc-index`, the embedded indexer follows
    /// the tip asynchronously and is not a blocker for checkpoint execution, so
    /// reads of the index surface must wait for it explicitly.
    async fn wait_for_rpc_index_on_handle(
        handle: &SuiNodeHandle,
        checkpoint: u64,
        wait_for_history: bool,
    ) {
        // Skip nodes without an embedded index; there is nothing to wait for.
        if handle.with(|node| node.embedded_rpc_store().is_none()) {
            return;
        }
        let deadline = Instant::now() + Duration::from_secs(60);
        loop {
            let (live_committed, history_committed) = handle.with(|node| {
                node.embedded_rpc_store()
                    .map(|embedded| {
                        (
                            embedded.live_committed_checkpoint(),
                            embedded.history_committed_checkpoint(),
                        )
                    })
                    .unwrap_or((None, None))
            });
            if live_committed.is_some_and(|c| c >= checkpoint)
                && (!wait_for_history || history_committed.is_some_and(|c| c >= checkpoint))
            {
                return;
            }
            assert!(
                Instant::now() < deadline,
                "timed out waiting for the embedded rpc-store to index checkpoint \
                 {checkpoint} (live committed = {live_committed:?}, \
                 history committed = {history_committed:?})",
            );
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    }

    /// Wait until the rpc fullnode's embedded rpc-store has indexed through its
    /// current highest executed checkpoint. Call after building the cluster so
    /// genesis data is queryable through index surfaces before tests issue
    /// their first index reads. No-op when the fullnode has indexing disabled.
    pub async fn wait_for_rpc_index_ready(&self) {
        let handle = &self.fullnode_handle.sui_node;
        let highest_executed = handle.with(|node| {
            node.state()
                .get_checkpoint_store()
                .get_highest_executed_checkpoint_seq_number()
                .expect("db error")
                .unwrap_or(0)
        });
        Self::wait_for_rpc_index_on_handle(handle, highest_executed, true).await;
    }

    /// Execute a transaction on the network and wait for it to be executed on the rpc fullnode.
    /// Also expects the effects status to be ExecutionStatus::Success.
    /// This function is recommended for transaction execution since it most resembles the
    /// production path.
    pub async fn execute_transaction(&self, tx: Transaction) -> ExecutedTransaction {
        self.wallet.execute_transaction_must_succeed(tx).await
    }

    /// Different from `execute_transaction` which returns RPC effects types, this function
    /// returns raw effects and events from the transaction driver.
    /// It also does not check whether the transaction is executed successfully.
    /// Before returning, it waits for the transaction to settle on the fullnode so that
    /// subsequent queries there read consistent results.
    pub async fn execute_transaction_return_raw_effects(
        &self,
        tx: Transaction,
    ) -> anyhow::Result<(TransactionEffects, TransactionEvents)> {
        let digest = *tx.digest();
        let results = self.submit_and_execute(tx, None).await?;
        self.wait_for_tx_settlement(&[digest]).await;
        Ok(results)
    }

    pub fn authority_aggregator(&self) -> Arc<AuthorityAggregator<NetworkAuthorityClient>> {
        self.fullnode_handle
            .sui_node
            .with(|node| node.clone_authority_aggregator().unwrap())
    }

    /// Submit a transaction through the transaction driver and wait for finality.
    /// Returns the raw transaction effects and events without checking execution status.
    pub async fn submit_and_execute(
        &self,
        tx: Transaction,
        client_addr: Option<SocketAddr>,
    ) -> anyhow::Result<(TransactionEffects, TransactionEvents)> {
        let transaction_driver = self.fullnode_handle.sui_node.with(|node| {
            node.transaction_orchestrator()
                .expect("fullnode must have a transaction orchestrator")
                .transaction_driver()
                .clone()
        });
        let response = transaction_driver
            .drive_transaction(
                SubmitTxRequest::new_transaction(tx),
                SubmitTransactionOptions {
                    forwarded_client_addr: client_addr,
                    ..Default::default()
                },
                Some(TRANSACTION_FINALITY_TIMEOUT),
            )
            .await?;

        Ok((
            response.effects.effects,
            response.events.unwrap_or_default(),
        ))
    }

    /// This call sends some funds from the seeded address to the funding
    /// address for the given amount and returns the gas object ref. This
    /// is useful to construct transactions from the funding address.
    pub async fn fund_address_and_return_gas(
        &self,
        rgp: u64,
        amount: Option<u64>,
        funding_address: SuiAddress,
    ) -> ObjectRef {
        let context = &self.wallet;
        let (sender, gas) = context.get_one_gas_object().await.unwrap().unwrap();
        let tx = context
            .sign_transaction(
                &TestTransactionBuilder::new(sender, gas, rgp)
                    .transfer_sui(amount, funding_address)
                    .build(),
            )
            .await;
        context.execute_transaction_must_succeed(tx).await;

        context
            .get_one_gas_object_owned_by_address(funding_address)
            .await
            .unwrap()
            .unwrap()
    }

    pub async fn transfer_sui_must_exceed(
        &self,
        sender: SuiAddress,
        receiver: SuiAddress,
        amount: u64,
    ) -> ObjectID {
        let tx = self
            .test_transaction_builder_with_sender(sender)
            .await
            .transfer_sui(Some(amount), receiver)
            .build();
        let effects = self.sign_and_execute_transaction(&tx).await.effects;
        assert!(effects.status().is_ok());
        effects.created().first().unwrap().0.0
    }

    #[cfg(msim)]
    pub fn set_safe_mode_expected(&self, value: bool) {
        for n in self.all_node_handles() {
            n.with(|node| node.set_safe_mode_expected(value));
        }
    }
}

pub struct RandomNodeRestarter {
    test_cluster: Arc<TestCluster>,

    // How frequently should we kill nodes
    kill_interval: Uniform<Duration>,
    // How long should we wait before restarting them.
    restart_delay: Uniform<Duration>,

    task_handle: Mutex<Option<JoinHandle<()>>>,
}

impl RandomNodeRestarter {
    fn new(test_cluster: Arc<TestCluster>) -> Self {
        Self {
            test_cluster,
            kill_interval: Uniform::new(Duration::from_secs(10), Duration::from_secs(11)),
            restart_delay: Uniform::new(Duration::from_secs(1), Duration::from_secs(2)),
            task_handle: Default::default(),
        }
    }

    pub fn with_kill_interval_secs(mut self, a: u64, b: u64) -> Self {
        self.kill_interval = Uniform::new(Duration::from_secs(a), Duration::from_secs(b));
        self
    }

    pub fn with_restart_delay_secs(mut self, a: u64, b: u64) -> Self {
        self.restart_delay = Uniform::new(Duration::from_secs(a), Duration::from_secs(b));
        self
    }

    pub fn run(&self) {
        let test_cluster = self.test_cluster.clone();
        let kill_interval = self.kill_interval;
        let restart_delay = self.restart_delay;
        let validators = self.test_cluster.get_validator_pubkeys();
        let mut task_handle = self.task_handle.lock().unwrap();
        assert!(task_handle.is_none());
        task_handle.replace(tokio::task::spawn(async move {
            loop {
                let delay = kill_interval.sample(&mut OsRng);
                info!("Sleeping {delay:?} before killing a validator");
                sleep(delay).await;

                let validator = validators.choose(&mut OsRng).unwrap();
                info!("Killing validator {:?}", validator.concise());
                test_cluster.stop_node(validator);

                let delay = restart_delay.sample(&mut OsRng);
                info!("Sleeping {delay:?} before restarting");
                sleep(delay).await;
                info!("Starting validator {:?}", validator.concise());
                test_cluster.start_node(validator).await;
            }
        }));
    }
}

impl Drop for RandomNodeRestarter {
    fn drop(&mut self) {
        if let Some(handle) = self.task_handle.lock().unwrap().take() {
            handle.abort();
        }
    }
}

pub struct TestClusterBuilder {
    genesis_config: Option<GenesisConfig>,
    network_config: Option<NetworkConfig>,
    additional_objects: Vec<Object>,
    num_validators: Option<usize>,
    validators: Option<Vec<ValidatorGenesisConfig>>,
    fullnode_rpc_port: Option<u16>,
    enable_fullnode_events: bool,
    disable_fullnode_pruning: bool,
    validator_supported_protocol_versions_config: ProtocolVersionsConfig,
    // Default to validator_supported_protocol_versions_config, but can be overridden.
    fullnode_supported_protocol_versions_config: Option<ProtocolVersionsConfig>,
    db_checkpoint_config_validators: DBCheckpointConfig,
    db_checkpoint_config_fullnodes: DBCheckpointConfig,
    num_unpruned_validators: Option<usize>,
    jwk_fetch_interval: Option<Duration>,
    config_dir: Option<PathBuf>,
    default_jwks: bool,
    authority_overload_config: Option<AuthorityOverloadConfig>,
    consensus_transaction_pool_config: Option<ConsensusTransactionPoolConfig>,
    execution_cache_config: Option<ExecutionCacheConfig>,
    data_ingestion_dir: Option<PathBuf>,
    fullnode_run_with_range: Option<RunWithRange>,
    fullnode_policy_config: Option<PolicyConfig>,
    fullnode_fw_config: Option<RemoteFirewallConfig>,

    validator_global_state_hash_v2_enabled_config: GlobalStateHashV2EnabledConfig,
    validator_funds_withdraw_scheduler_type_config: FundsWithdrawSchedulerTypeConfig,

    rpc_config: Option<sui_config::RpcConfig>,

    chain_override: Option<Chain>,

    execution_time_observer_config: Option<sui_config::node::ExecutionTimeObserverConfig>,

    validator_observer_config: Option<ValidatorObserverConfigCallback>,

    observer_fullnode: bool,

    state_sync_config: Option<sui_config::p2p::StateSyncConfig>,

    peer_deny_sync_config_callback:
        Option<sui_swarm_config::network_config_builder::PeerDenySyncConfigCallback>,

    #[cfg(msim)]
    inject_synthetic_execution_time: bool,
}

impl TestClusterBuilder {
    pub fn new() -> Self {
        TestClusterBuilder {
            genesis_config: None,
            network_config: None,
            chain_override: None,
            additional_objects: vec![],
            fullnode_rpc_port: None,
            num_validators: None,
            validators: None,
            enable_fullnode_events: false,
            disable_fullnode_pruning: false,
            validator_supported_protocol_versions_config: ProtocolVersionsConfig::Default,
            fullnode_supported_protocol_versions_config: None,
            db_checkpoint_config_validators: DBCheckpointConfig::default(),
            db_checkpoint_config_fullnodes: DBCheckpointConfig::default(),
            num_unpruned_validators: None,
            jwk_fetch_interval: None,
            config_dir: None,
            default_jwks: false,
            authority_overload_config: None,
            consensus_transaction_pool_config: None,
            execution_cache_config: None,
            data_ingestion_dir: None,
            fullnode_run_with_range: None,
            fullnode_policy_config: None,
            fullnode_fw_config: None,
            validator_global_state_hash_v2_enabled_config: GlobalStateHashV2EnabledConfig::Global(
                true,
            ),
            validator_funds_withdraw_scheduler_type_config:
                FundsWithdrawSchedulerTypeConfig::PerValidator(Arc::new(|idx| {
                    if idx % 2 == 0 {
                        FundsWithdrawSchedulerType::Eager
                    } else {
                        FundsWithdrawSchedulerType::Naive
                    }
                })),
            rpc_config: None,
            execution_time_observer_config: None,
            validator_observer_config: None,
            observer_fullnode: false,
            state_sync_config: None,
            peer_deny_sync_config_callback: None,
            #[cfg(msim)]
            inject_synthetic_execution_time: false,
        }
    }

    pub fn with_state_sync_config(mut self, config: sui_config::p2p::StateSyncConfig) -> Self {
        self.state_sync_config = Some(config);
        self
    }

    /// Per-validator hook for `peer_deny_sync_config`. The closure receives this
    /// validator's authority name and the slice of all genesis-committee authority
    /// names, so callers can compute an allowlist that references peers (e.g.
    /// "trust everyone but myself").
    pub fn with_peer_deny_sync_config_per_validator(
        mut self,
        f: sui_swarm_config::network_config_builder::PeerDenySyncConfigCallback,
    ) -> Self {
        self.peer_deny_sync_config_callback = Some(f);
        self
    }

    pub fn with_execution_time_observer_config(
        mut self,
        config: sui_config::node::ExecutionTimeObserverConfig,
    ) -> Self {
        self.execution_time_observer_config = Some(config);
        self
    }

    pub fn with_validator_observer_config(mut self, c: ValidatorObserverConfigCallback) -> Self {
        self.validator_observer_config = Some(c);
        self
    }

    /// Adds a fullnode to the cluster that syncs as a consensus observer, subscribed to the
    /// first validator. Unless a validator observer config is provided, the first validator
    /// gets its observer server enabled. The observer is available via
    /// `TestCluster::observer_node()`.
    ///
    /// The first validator must end up with its observer server enabled, which `build()`
    /// validates: a callback passed to `with_validator_observer_config` must return
    /// `Some(_)` for index 0, and a prebuilt network config passed to `set_network_config`
    /// must already enable the observer server on the first validator (the callback is not
    /// applied to prebuilt network configs).
    pub fn with_observer_fullnode(mut self) -> Self {
        self.observer_fullnode = true;
        self
    }

    pub fn with_fullnode_run_with_range(mut self, run_with_range: Option<RunWithRange>) -> Self {
        if let Some(run_with_range) = run_with_range {
            self.fullnode_run_with_range = Some(run_with_range);
        }
        self
    }

    pub fn with_fullnode_policy_config(mut self, config: Option<PolicyConfig>) -> Self {
        self.fullnode_policy_config = config;
        self
    }

    pub fn with_fullnode_fw_config(mut self, config: Option<RemoteFirewallConfig>) -> Self {
        self.fullnode_fw_config = config;
        self
    }

    pub fn with_fullnode_rpc_port(mut self, rpc_port: u16) -> Self {
        self.fullnode_rpc_port = Some(rpc_port);
        self
    }

    pub fn set_genesis_config(mut self, genesis_config: GenesisConfig) -> Self {
        assert!(self.genesis_config.is_none() && self.network_config.is_none());
        self.genesis_config = Some(genesis_config);
        self
    }

    pub fn set_network_config(mut self, network_config: NetworkConfig) -> Self {
        assert!(self.genesis_config.is_none() && self.network_config.is_none());
        self.network_config = Some(network_config);
        self
    }

    pub fn with_objects<I: IntoIterator<Item = Object>>(mut self, objects: I) -> Self {
        self.additional_objects.extend(objects);
        self
    }

    /// Set the number of default validators to spawn. Can be overridden by `with_validators`, if
    /// you need to provide more specific genesis configs for each validator.
    pub fn with_num_validators(mut self, num: usize) -> Self {
        self.num_validators = Some(num);
        self
    }

    /// Provide validator genesis configs, overrides the `num_validators` setting.
    pub fn with_validators(mut self, validators: Vec<ValidatorGenesisConfig>) -> Self {
        self.validators = Some(validators);
        self
    }

    pub fn enable_fullnode_events(mut self) -> Self {
        self.enable_fullnode_events = true;
        self
    }

    pub fn disable_fullnode_pruning(mut self) -> Self {
        self.disable_fullnode_pruning = true;
        self
    }

    pub fn with_enable_db_checkpoints_validators(mut self) -> Self {
        self.db_checkpoint_config_validators = DBCheckpointConfig {
            perform_db_checkpoints_at_epoch_end: true,
            checkpoint_path: None,
            object_store_config: None,
            perform_index_db_checkpoints_at_epoch_end: None,
            prune_and_compact_before_upload: None,
        };
        self
    }

    pub fn with_enable_db_checkpoints_fullnodes(mut self) -> Self {
        self.db_checkpoint_config_fullnodes = DBCheckpointConfig {
            perform_db_checkpoints_at_epoch_end: true,
            checkpoint_path: None,
            object_store_config: None,
            perform_index_db_checkpoints_at_epoch_end: None,
            prune_and_compact_before_upload: Some(true),
        };
        self
    }

    pub fn with_epoch_duration_ms(mut self, epoch_duration_ms: u64) -> Self {
        assert!(
            epoch_duration_ms >= 10000,
            "Epoch duration must be at least 10s (10000ms) to avoid flaky tests. Got {epoch_duration_ms}ms."
        );
        self.get_or_init_genesis_config()
            .parameters
            .epoch_duration_ms = epoch_duration_ms;
        self
    }

    pub fn with_stake_subsidy_start_epoch(mut self, stake_subsidy_start_epoch: u64) -> Self {
        self.get_or_init_genesis_config()
            .parameters
            .stake_subsidy_start_epoch = stake_subsidy_start_epoch;
        self
    }

    pub fn with_supported_protocol_versions(mut self, c: SupportedProtocolVersions) -> Self {
        self.validator_supported_protocol_versions_config = ProtocolVersionsConfig::Global(c);
        self
    }

    pub fn with_jwk_fetch_interval(mut self, i: Duration) -> Self {
        self.jwk_fetch_interval = Some(i);
        self
    }

    pub fn with_fullnode_supported_protocol_versions_config(
        mut self,
        c: SupportedProtocolVersions,
    ) -> Self {
        self.fullnode_supported_protocol_versions_config = Some(ProtocolVersionsConfig::Global(c));
        self
    }

    pub fn with_protocol_version(mut self, v: ProtocolVersion) -> Self {
        self.get_or_init_genesis_config()
            .parameters
            .protocol_version = v;
        self
    }

    pub fn with_supported_protocol_version_callback(
        mut self,
        func: SupportedProtocolVersionsCallback,
    ) -> Self {
        self.validator_supported_protocol_versions_config =
            ProtocolVersionsConfig::PerValidator(func);
        self
    }

    pub fn with_global_state_hash_v2_enabled_callback(
        mut self,
        func: GlobalStateHashV2EnabledCallback,
    ) -> Self {
        self.validator_global_state_hash_v2_enabled_config =
            GlobalStateHashV2EnabledConfig::PerValidator(func);
        self
    }

    pub fn with_validator_candidates(
        mut self,
        addresses: impl IntoIterator<Item = SuiAddress>,
    ) -> Self {
        self.get_or_init_genesis_config()
            .accounts
            .extend(addresses.into_iter().map(|address| AccountConfig {
                address: Some(address),
                gas_amounts: vec![DEFAULT_GAS_AMOUNT, DEFAULT_GAS_AMOUNT],
            }));
        self
    }

    pub fn with_num_unpruned_validators(mut self, n: usize) -> Self {
        self.num_unpruned_validators = Some(n);
        self
    }

    pub fn with_accounts(mut self, accounts: Vec<AccountConfig>) -> Self {
        self.get_or_init_genesis_config().accounts = accounts;
        self
    }

    pub fn with_additional_accounts(mut self, accounts: Vec<AccountConfig>) -> Self {
        self.get_or_init_genesis_config().accounts.extend(accounts);
        self
    }

    pub fn with_config_dir(mut self, config_dir: PathBuf) -> Self {
        self.config_dir = Some(config_dir);
        self
    }

    pub fn with_default_jwks(mut self) -> Self {
        self.default_jwks = true;
        self
    }

    pub fn with_authority_overload_config(mut self, config: AuthorityOverloadConfig) -> Self {
        assert!(self.network_config.is_none());
        self.authority_overload_config = Some(config);
        self
    }

    pub fn with_consensus_transaction_pool_config(
        mut self,
        config: ConsensusTransactionPoolConfig,
    ) -> Self {
        assert!(self.network_config.is_none());
        self.consensus_transaction_pool_config = Some(config);
        self
    }

    pub fn with_execution_cache_config(mut self, config: ExecutionCacheConfig) -> Self {
        assert!(self.network_config.is_none());
        self.execution_cache_config = Some(config);
        self
    }

    pub fn with_data_ingestion_dir(mut self, path: PathBuf) -> Self {
        self.data_ingestion_dir = Some(path);
        self
    }

    pub fn with_rpc_config(mut self, config: sui_config::RpcConfig) -> Self {
        self.rpc_config = Some(config);
        self
    }

    pub fn with_chain_override(mut self, chain: Chain) -> Self {
        self.chain_override = Some(chain);
        self
    }

    #[cfg(msim)]
    pub fn with_synthetic_execution_time_injection(mut self) -> Self {
        self.inject_synthetic_execution_time = true;
        self
    }

    pub async fn build(mut self) -> TestCluster {
        // All test clusters receive a continuous stream of random JWKs.
        // If we later use zklogin authenticated transactions in tests we will need to supply
        // valid JWKs as well.
        #[cfg(msim)]
        if !self.default_jwks {
            sui_node::set_jwk_injector(Arc::new(|_authority, provider| {
                use fastcrypto_zkp::bn254::zk_login::{JWK, JwkId};
                use rand::Rng;

                // generate random (and possibly conflicting) id/key pairings.
                let id_num = rand::thread_rng().gen_range(1..=4);
                let key_num = rand::thread_rng().gen_range(1..=4);

                let id = JwkId {
                    iss: provider.get_config().iss,
                    kid: format!("kid{}", id_num),
                };

                let jwk = JWK {
                    kty: "kty".to_string(),
                    e: "e".to_string(),
                    n: format!("n{}", key_num),
                    alg: "alg".to_string(),
                };

                Ok(vec![(id, jwk)])
            }));
        }

        if self.observer_fullnode {
            // The observer fullnode subscribes to the first validator, so its observer server
            // must be enabled. Validate this up front to fail with an actionable error instead
            // of a deep panic in `observer_peer_record` when the observer's config is built.
            if let Some(network_config) = &self.network_config {
                // A prebuilt network config bypasses the validator observer config callback,
                // so it must already have the observer server enabled.
                let observer_enabled = network_config
                    .validator_configs()
                    .first()
                    .and_then(|c| c.consensus_config())
                    .and_then(|c| c.parameters.as_ref())
                    .and_then(|p| p.observer.server_port)
                    .is_some();
                assert!(
                    observer_enabled,
                    "with_observer_fullnode() requires the first validator's observer server \
                     to be enabled, but the network config passed to set_network_config() does \
                     not enable it. Enable the observer server on the first validator when \
                     building the network config."
                );
            } else if let Some(cb) = &self.validator_observer_config {
                assert!(
                    cb(0).is_some(),
                    "with_observer_fullnode() requires the validator observer config callback \
                     to enable the observer server on the first validator (index 0)."
                );
            } else {
                self.validator_observer_config = Some(Arc::new(|idx| {
                    (idx == 0).then(consensus_config::ObserverParameters::default)
                }));
            }
        }

        let mut swarm = self.start_swarm().await.unwrap();
        let working_dir = swarm.dir().to_path_buf();

        let fullnode = swarm.fullnodes().next().unwrap();
        let json_rpc_address = fullnode.config().json_rpc_address;
        let fullnode_handle =
            FullNodeHandle::new(fullnode.get_node_handle().unwrap(), json_rpc_address).await;

        if self.observer_fullnode {
            // Boxed so the frame (NodeConfig by value, held across the await) lives on
            // the heap instead of inflating build()'s state machine for every caller,
            // most of which never enable the observer. Unoptimized async frames are
            // large enough that this tips borderline test binaries over the 2 MiB
            // tokio worker stack.
            Box::pin(async {
                let mut config = swarm
                    .get_fullnode_config_builder()
                    .with_observer_subscribed_to_validator(0)
                    .build(&mut OsRng, swarm.config());
                // An observer that halts at a checkpoint range boundary defeats its purpose.
                config.run_with_range = None;
                // Deliberately drop the returned node handle: holding it would keep the
                // instance alive across a crash and prevent the simulator from restarting
                // the node. Access the observer via `TestCluster::observer_node()`.
                swarm.spawn_new_node(config).await;
            })
            .await;
        }

        let mut wallet_conf: SuiClientConfig =
            PersistedConfig::read(&working_dir.join(SUI_CLIENT_CONFIG)).unwrap();
        wallet_conf.envs.push(SuiEnv {
            alias: "localnet".to_string(),
            rpc: fullnode_handle.rpc_url.clone(),
            ws: None,
            basic_auth: None,
            chain_id: None,
        });
        wallet_conf.active_env = Some("localnet".to_string());

        wallet_conf
            .persisted(&working_dir.join(SUI_CLIENT_CONFIG))
            .save()
            .unwrap();

        let wallet_conf = swarm.dir().join(SUI_CLIENT_CONFIG);
        let wallet = WalletContext::new(&wallet_conf).unwrap();

        let cluster = TestCluster {
            swarm,
            wallet,
            fullnode_handle,
        };

        // The embedded rpc-store indexes the tip asynchronously, so genesis
        // data is not queryable through every index surface the instant the node
        // is up. Wait for it before handing the cluster to tests.
        cluster.wait_for_rpc_index_ready().await;

        cluster
    }

    /// Start a Swarm and set up WalletConfig
    async fn start_swarm(&mut self) -> Result<Swarm, anyhow::Error> {
        let mut builder: SwarmBuilder = Swarm::builder()
            .with_objects(self.additional_objects.clone())
            .with_db_checkpoint_config(self.db_checkpoint_config_validators.clone())
            .with_supported_protocol_versions_config(
                self.validator_supported_protocol_versions_config.clone(),
            )
            .with_global_state_hash_v2_enabled_config(
                self.validator_global_state_hash_v2_enabled_config.clone(),
            )
            .with_funds_withdraw_scheduler_type_config(
                self.validator_funds_withdraw_scheduler_type_config.clone(),
            )
            .with_fullnode_count(1)
            .with_fullnode_supported_protocol_versions_config(
                self.fullnode_supported_protocol_versions_config
                    .clone()
                    .unwrap_or(self.validator_supported_protocol_versions_config.clone()),
            )
            .with_db_checkpoint_config(self.db_checkpoint_config_fullnodes.clone())
            .with_fullnode_run_with_range(self.fullnode_run_with_range)
            .with_fullnode_policy_config(self.fullnode_policy_config.clone())
            .with_fullnode_fw_config(self.fullnode_fw_config.clone());

        if let Some(validators) = self.validators.take() {
            builder = builder.with_validators(validators);
        } else {
            builder = builder.committee_size(
                NonZeroUsize::new(self.num_validators.unwrap_or(NUM_VALIDATOR)).unwrap(),
            )
        };

        if let Some(chain) = self.chain_override {
            builder = builder.with_chain_override(chain);
        }

        if let Some(genesis_config) = self.genesis_config.take() {
            builder = builder.with_genesis_config(genesis_config);
        }

        if let Some(network_config) = self.network_config.take() {
            builder = builder.with_network_config(network_config);
        }

        if let Some(authority_overload_config) = self.authority_overload_config.take() {
            builder = builder.with_authority_overload_config(authority_overload_config);
        }

        if let Some(config) = self.consensus_transaction_pool_config.take() {
            builder = builder.with_consensus_transaction_pool_config(config);
        }

        if let Some(execution_cache_config) = self.execution_cache_config.take() {
            builder = builder.with_execution_cache_config(execution_cache_config);
        }

        if let Some(fullnode_rpc_port) = self.fullnode_rpc_port {
            builder = builder.with_fullnode_rpc_port(fullnode_rpc_port);
        }

        if let Some(rpc_config) = &self.rpc_config {
            builder = builder.with_fullnode_rpc_config(rpc_config.clone());
        }
        if let Some(num_unpruned_validators) = self.num_unpruned_validators {
            builder = builder.with_num_unpruned_validators(num_unpruned_validators);
        }

        if let Some(jwk_fetch_interval) = self.jwk_fetch_interval {
            builder = builder.with_jwk_fetch_interval(jwk_fetch_interval);
        }

        if let Some(config_dir) = self.config_dir.take() {
            builder = builder.dir(config_dir);
        }

        if let Some(data_ingestion_dir) = self.data_ingestion_dir.take() {
            builder = builder.with_data_ingestion_dir(data_ingestion_dir);
        }

        if let Some(state_sync_config) = self.state_sync_config.clone() {
            builder = builder.with_state_sync_config(state_sync_config);
        }

        if let Some(cb) = self.peer_deny_sync_config_callback.clone() {
            builder = builder.with_peer_deny_sync_config_per_validator(cb);
        }

        if self.disable_fullnode_pruning {
            builder = builder.with_disable_fullnode_pruning();
        }

        if let Some(validator_observer_config) = self.validator_observer_config.take() {
            builder = builder.with_validator_observer_config(validator_observer_config);
        }

        #[cfg(msim)]
        {
            if let Some(mut config) = self.execution_time_observer_config.clone() {
                if self.inject_synthetic_execution_time {
                    config.inject_synthetic_execution_time = Some(true);
                }
                builder = builder.with_execution_time_observer_config(config);
            } else if self.inject_synthetic_execution_time {
                use sui_config::node::ExecutionTimeObserverConfig;

                let mut config = ExecutionTimeObserverConfig::default();
                config.inject_synthetic_execution_time = Some(true);
                builder = builder.with_execution_time_observer_config(config);
            }
        }

        let mut swarm = builder.build();
        swarm.launch().await?;

        let dir = swarm.dir();

        let network_path = dir.join(SUI_NETWORK_CONFIG);
        let wallet_path = dir.join(SUI_CLIENT_CONFIG);
        let keystore_path = dir.join(SUI_KEYSTORE_FILENAME);

        swarm.config().save(network_path)?;
        let mut keystore = Keystore::from(FileBasedKeystore::load_or_create(&keystore_path)?);
        for key in &swarm.config().account_keys {
            keystore
                .import(None, SuiKeyPair::Ed25519(key.copy()))
                .await?;
        }

        let active_address = keystore.addresses().first().cloned();

        // Create wallet config with stated authorities port
        SuiClientConfig {
            keystore: Keystore::from(FileBasedKeystore::load_or_create(&keystore_path)?),
            external_keys: None,
            envs: Default::default(),
            active_address,
            active_env: Default::default(),
        }
        .save(wallet_path)?;

        // Return network handle
        Ok(swarm)
    }

    fn get_or_init_genesis_config(&mut self) -> &mut GenesisConfig {
        if self.genesis_config.is_none() {
            self.genesis_config = Some(GenesisConfig::for_local_testing());
        }
        self.genesis_config.as_mut().unwrap()
    }
}

impl Default for TestClusterBuilder {
    fn default() -> Self {
        Self::new()
    }
}