iota_cluster_test/

cluster.rs

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

use std::{net::SocketAddr, path::Path};

use async_trait::async_trait;
use iota_config::{Config, IOTA_GENESIS_FILENAME, IOTA_KEYSTORE_FILENAME, node::Genesis};
use iota_graphql_rpc::{
    config::ConnectionConfig, test_infra::cluster::start_graphql_server_with_fn_rpc,
};
use iota_indexer::test_utils::{IndexerTypeConfig, start_test_indexer};
use iota_keys::keystore::{AccountKeystore, FileBasedKeystore, Keystore};
use iota_sdk::{
    iota_client_config::{IotaClientConfig, IotaEnv},
    wallet_context::WalletContext,
};
use iota_swarm::memory::Swarm;
use iota_swarm_config::{genesis_config::GenesisConfig, network_config::PersistedNetworkConfig};
use iota_types::crypto::AccountPrivateKey;
use tempfile::tempdir;
use test_cluster::{TestCluster, TestClusterBuilder};
use tracing::info;

use super::config::{ClusterTestOpt, Env};

const DEVNET_FAUCET_ADDR: &str = "https://faucet.devnet.iota.cafe:443";
const TESTNET_FAUCET_ADDR: &str = "https://faucet.testnet.iota.cafe:443";
const DEVNET_FULLNODE_ADDR: &str = "https://api.devnet.iota.cafe:443";
const TESTNET_FULLNODE_ADDR: &str = "https://api.testnet.iota.cafe:443";

pub struct ClusterFactory;

impl ClusterFactory {
    pub async fn start(
        options: &ClusterTestOpt,
    ) -> Result<Box<dyn Cluster + Sync + Send>, anyhow::Error> {
        Ok(match &options.env {
            Env::NewLocal => Box::new(LocalNewCluster::start(options).await?) as Box<_>,
            _ => Box::new(RemoteRunningCluster::start(options).await?) as Box<_>,
        })
    }
}

/// Cluster Abstraction
#[async_trait]
pub trait Cluster {
    async fn start(options: &ClusterTestOpt) -> Result<Self, anyhow::Error>
    where
        Self: Sized;

    fn fullnode_url(&self) -> &str;

    /// gRPC endpoint of the fullnode, when one is available (local clusters).
    fn grpc_url(&self) -> Option<&str>;

    fn user_key(&self) -> AccountPrivateKey;
    fn indexer_url(&self) -> &Option<String>;

    /// Returns faucet url in a remote cluster.
    fn remote_faucet_url(&self) -> Option<&str>;

    /// Returns faucet key in a local cluster.
    fn local_faucet_key(&self) -> Option<&AccountPrivateKey>;

    /// Place to put config for the wallet, and any locally running services.
    fn config_directory(&self) -> &Path;
}

/// Represents an up and running cluster deployed remotely.
pub struct RemoteRunningCluster {
    fullnode_url: String,
    faucet_url: String,
    config_directory: tempfile::TempDir,
}

#[async_trait]
impl Cluster for RemoteRunningCluster {
    async fn start(options: &ClusterTestOpt) -> Result<Self, anyhow::Error> {
        let (fullnode_url, faucet_url) = match options.env {
            Env::Devnet => (
                String::from(DEVNET_FULLNODE_ADDR),
                String::from(DEVNET_FAUCET_ADDR),
            ),
            Env::Testnet => (
                String::from(TESTNET_FULLNODE_ADDR),
                String::from(TESTNET_FAUCET_ADDR),
            ),
            Env::CustomRemote => (
                options
                    .fullnode_address
                    .clone()
                    .expect("expect 'fullnode_address' for Env::Custom"),
                options
                    .faucet_address
                    .clone()
                    .expect("expect 'faucet_address' for Env::Custom"),
            ),
            Env::NewLocal => {
                unreachable!("the NewLocal variant shouldn't use RemoteRunningCluster")
            }
        };

        // TODO: test connectivity before proceeding?

        Ok(Self {
            fullnode_url,
            faucet_url,
            config_directory: tempfile::tempdir()?,
        })
    }

    fn fullnode_url(&self) -> &str {
        &self.fullnode_url
    }

    fn grpc_url(&self) -> Option<&str> {
        None
    }

    fn indexer_url(&self) -> &Option<String> {
        &None
    }

    fn user_key(&self) -> AccountPrivateKey {
        AccountPrivateKey::random()
    }

    fn remote_faucet_url(&self) -> Option<&str> {
        Some(&self.faucet_url)
    }

    fn local_faucet_key(&self) -> Option<&AccountPrivateKey> {
        None
    }

    fn config_directory(&self) -> &Path {
        self.config_directory.path()
    }
}

/// Represents a local Cluster which starts per cluster test run.
pub struct LocalNewCluster {
    test_cluster: TestCluster,
    fullnode_url: String,
    grpc_url: String,
    indexer_url: Option<String>,
    faucet_key: AccountPrivateKey,
    config_directory: tempfile::TempDir,
}

impl LocalNewCluster {
    #[allow(unused)]
    pub fn swarm(&self) -> &Swarm {
        &self.test_cluster.swarm
    }
}

#[async_trait]
impl Cluster for LocalNewCluster {
    async fn start(options: &ClusterTestOpt) -> Result<Self, anyhow::Error> {
        let data_ingestion_path = tempdir()?.keep();
        // TODO: options should contain port instead of address
        let fullnode_rpc_addr = options.fullnode_address.as_ref().map(|addr| {
            addr.parse::<SocketAddr>()
                .expect("unable to parse fullnode address")
        });

        let indexer_address = options.indexer_address.as_ref().map(|addr| {
            addr.parse::<SocketAddr>()
                .expect("unable to parse indexer address")
        });

        let mut cluster_builder = TestClusterBuilder::new()
            .enable_fullnode_events()
            .with_data_ingestion_dir(data_ingestion_path.clone())
            // disable full node pruning: tests read `balance_changes` /
            // `object_changes` which needs historical data.
            .disable_fullnode_pruning();

        // Check if we already have a config directory that is passed
        if let Some(config_dir) = options.config_dir.clone() {
            assert!(options.epoch_duration_ms.is_none());
            // Derive the node configs of the network from its persisted state.
            let network_config = PersistedNetworkConfig::read(&config_dir)?.into_network_config(
                &config_dir,
                Genesis::new_from_file(config_dir.join(IOTA_GENESIS_FILENAME)),
            )?;
            cluster_builder = cluster_builder.set_network_config(network_config);

            cluster_builder = cluster_builder.with_config_dir(config_dir);
        } else {
            // Let the faucet account hold 1000 gas objects on genesis
            let genesis_config = GenesisConfig::custom_genesis(1, 100);
            // Custom genesis should be build here where we add the extra accounts
            cluster_builder = cluster_builder.set_genesis_config(genesis_config);

            if let Some(epoch_duration_ms) = options.epoch_duration_ms {
                cluster_builder = cluster_builder.with_epoch_duration_ms(epoch_duration_ms);
            }
        }

        if let Some(fullnode_rpc_addr) = fullnode_rpc_addr {
            cluster_builder = cluster_builder.with_fullnode_rpc_addr(fullnode_rpc_addr);
        }

        let mut test_cluster = cluster_builder.build().await;

        // Use the wealthy account for faucet
        let faucet_key = test_cluster.swarm.config_mut().account_keys.swap_remove(0);
        let faucet_address = faucet_key.public_key().derive_address();
        info!(?faucet_address, "faucet_address");

        // This cluster has fullnode handle, safe to unwrap
        let fullnode_url = test_cluster.fullnode_handle.rpc_url.clone();
        let grpc_url = test_cluster.grpc_url();

        if let (Some(pg_address), Some(indexer_address)) =
            (options.pg_address.clone(), indexer_address)
        {
            // Start in writer mode
            start_test_indexer(
                pg_address.clone(),
                // reset the existing db
                true,
                None,
                test_cluster.grpc_url(),
                IndexerTypeConfig::writer_mode(None),
                Some(data_ingestion_path.clone()),
            )
            .await;

            // Start in reader mode
            start_test_indexer(
                pg_address,
                false,
                None,
                test_cluster.grpc_url(),
                IndexerTypeConfig::reader_mode(indexer_address.to_string()),
                Some(data_ingestion_path),
            )
            .await;
        }

        if let Some(graphql_address) = &options.graphql_address {
            let graphql_address = graphql_address.parse::<SocketAddr>()?;
            let graphql_connection_config = ConnectionConfig::new(
                Some(graphql_address.port()),
                Some(graphql_address.ip().to_string()),
                options.pg_address.clone(),
                None,
                None,
                None,
                None,
                None,
            );

            start_graphql_server_with_fn_rpc(
                graphql_connection_config.clone(),
                Some(test_cluster.grpc_url()),
                // resolves to default cancellation_token
                None,
                // resolves to default service config
                None,
            )
            .await;
        }

        // Let nodes connect to one another
        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;

        // TODO: test connectivity before proceeding?
        Ok(Self {
            test_cluster,
            fullnode_url,
            grpc_url,
            faucet_key,
            config_directory: tempfile::tempdir()?,
            indexer_url: options.indexer_address.clone(),
        })
    }

    fn fullnode_url(&self) -> &str {
        &self.fullnode_url
    }

    fn grpc_url(&self) -> Option<&str> {
        Some(&self.grpc_url)
    }

    fn indexer_url(&self) -> &Option<String> {
        &self.indexer_url
    }

    fn user_key(&self) -> AccountPrivateKey {
        AccountPrivateKey::random()
    }

    fn remote_faucet_url(&self) -> Option<&str> {
        None
    }

    fn local_faucet_key(&self) -> Option<&AccountPrivateKey> {
        Some(&self.faucet_key)
    }

    fn config_directory(&self) -> &Path {
        self.config_directory.path()
    }
}

// Make linter happy
#[async_trait]
impl Cluster for Box<dyn Cluster + Send + Sync> {
    async fn start(_options: &ClusterTestOpt) -> Result<Self, anyhow::Error> {
        unreachable!(
            "if we already have a boxed Cluster trait object we wouldn't have to call this function"
        );
    }
    fn fullnode_url(&self) -> &str {
        (**self).fullnode_url()
    }
    fn grpc_url(&self) -> Option<&str> {
        (**self).grpc_url()
    }
    fn indexer_url(&self) -> &Option<String> {
        (**self).indexer_url()
    }

    fn user_key(&self) -> AccountPrivateKey {
        (**self).user_key()
    }

    fn remote_faucet_url(&self) -> Option<&str> {
        (**self).remote_faucet_url()
    }

    fn local_faucet_key(&self) -> Option<&AccountPrivateKey> {
        (**self).local_faucet_key()
    }

    fn config_directory(&self) -> &Path {
        (**self).config_directory()
    }
}

pub fn new_wallet_context_from_cluster(
    cluster: &(dyn Cluster + Sync + Send),
    private_key: AccountPrivateKey,
) -> WalletContext {
    let config_dir = cluster.config_directory();
    let wallet_config_path = config_dir.join("client.yaml");
    let fullnode_url = cluster.fullnode_url();
    info!("Use RPC: {fullnode_url}");
    let keystore_path = config_dir.join(IOTA_KEYSTORE_FILENAME);
    let mut keystore = Keystore::from(FileBasedKeystore::new(&keystore_path).unwrap());
    let address = private_key.public_key().derive_address();
    keystore.add_key(None, private_key).unwrap();
    IotaClientConfig::new(keystore)
        .with_envs([IotaEnv::new("localnet", fullnode_url)])
        .with_active_address(address)
        .with_active_env("localnet".to_string())
        .persisted(&wallet_config_path)
        .save()
        .unwrap();

    info!(
        "Initialize wallet from config path: {:?}",
        wallet_config_path
    );

    WalletContext::new(&wallet_config_path).unwrap_or_else(|e| {
        panic!("failed to init wallet context from path {wallet_config_path:?}, error: {e}")
    })
}