<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::smoke_test_environment::SmokeTestEnvironment;
use cli::client_proxy::ClientProxy;
use diem_config::config::{Identity, NodeConfig, SecureBackend};
use diem_crypto::ed25519::Ed25519PublicKey;
use diem_types::account_address::AccountAddress;
use rust_decimal::{prelude::FromPrimitive, Decimal};
use std::{collections::BTreeMap, fs::File, io::Write, path::PathBuf, str::FromStr};

// TODO(joshlind): Refactor all of these so that they can be contained within the calling
// test files and not shared across all tests.
pub fn compare_balances(
    expected_balances: Vec<(f64, String)>,
    extracted_balances: Vec<String>,
) -> bool {
    if extracted_balances.len() != extracted_balances.len() {
        return false;
    }

    let extracted_balances_dec: BTreeMap<_, _> = extracted_balances
        .into_iter()
        .map(|balance_str| {
            let (currency_code, stripped_str) = if balance_str.ends_with("XUS") {
                ("XUS", balance_str.trim_end_matches("XUS"))
            } else if balance_str.ends_with("XDX") {
                ("XDX", balance_str.trim_end_matches("XDX"))
            } else {
                panic!("Unexpected currency type returned for balance")
            };
            (currency_code, Decimal::from_str(stripped_str).ok())
        })
        .collect();

    expected_balances
        .into_iter()
        .all(|(balance, currency_code)| {
            if let Some(extracted_balance) = extracted_balances_dec.get(currency_code.as_str()) {
                Decimal::from_f64(balance) == *extracted_balance
            } else {
                false
            }
        })
}

/// Sets up a SmokeTestEnvironment with specified size and connects a client
/// proxy to the node_index.
pub fn setup_swarm_and_client_proxy(
    num_nodes: usize,
    node_index: usize,
) -> (SmokeTestEnvironment, ClientProxy) {
    let mut env = SmokeTestEnvironment::new(num_nodes);
    env.validator_swarm.launch();

    let client = env.get_validator_client(node_index, None);
    (env, client)
}

/// Waits for a transaction to be processed by all validator nodes in the smoke
/// test environment.
pub fn wait_for_transaction_on_all_nodes(
    env: &SmokeTestEnvironment,
    num_nodes: usize,
    account: AccountAddress,
    sequence_number: u64,
) {
    for i in 0..num_nodes {
        let client = env.get_validator_client(i, None);
        client
            .wait_for_transaction(account, sequence_number)
            .unwrap();
    }
}

/// This module provides useful functions for operating, handling and managing
/// DiemSwarm instances. It is particularly useful for working with tests that
/// require a SmokeTestEnvironment, as it provides a generic interface across
/// DiemSwarms, regardless of if the swarm is a validator swarm, validator full
/// node swarm, or a public full node swarm.
pub mod diem_swarm_utils {
    use crate::test_utils::fetch_backend_storage;
    use cli::client_proxy::ClientProxy;
    use diem_config::config::{NodeConfig, SecureBackend, WaypointConfig};
    use diem_events_fetcher::DiemEventsFetcher;
    use diem_key_manager::diem_interface::JsonRpcDiemInterface;
    use diem_operational_tool::test_helper::OperationalTool;
    use diem_secure_storage::{KVStorage, Storage};
    use diem_swarm::swarm::DiemSwarm;
    use diem_transaction_replay::DiemDebugger;
    use diem_types::{chain_id::ChainId, waypoint::Waypoint};
    use std::path::PathBuf;

    /// Returns a new client proxy connected to the given swarm at the specified
    /// node index.
    pub fn get_client_proxy(
        swarm: &DiemSwarm,
        node_index: usize,
        diem_root_key_path: &str,
        mnemonic_file_path: PathBuf,
        waypoint: Option<Waypoint>,
    ) -> ClientProxy {
        let port = swarm.get_client_port(node_index);

        let mnemonic_file_path = mnemonic_file_path
            .canonicalize()
            .expect("Unable to get canonical path of mnemonic_file_path")
            .to_str()
            .unwrap()
            .to_string();

        ClientProxy::new(
            ChainId::test(),
            &format!("http://localhost:{}/v1", port),
            &diem_root_key_path,
            &diem_root_key_path,
            &diem_root_key_path,
            false,
            /* faucet server */ None,
            Some(mnemonic_file_path),
            waypoint.unwrap_or(swarm.config.waypoint),
            true,
        )
        .unwrap()
    }

    /// Returns the JSON RPC url pointing to a node at the given
    /// node index.
    pub fn get_json_rpc_url(swarm: &DiemSwarm, node_index: usize) -> String {
        format!("http://127.0.0.1:{}", swarm.get_client_port(node_index))
    }

    /// Returns a JSON RPC based Diem Interface pointing to a node at the given
    /// node index.
    pub fn get_json_rpc_diem_interface(
        swarm: &DiemSwarm,
        node_index: usize,
    ) -> JsonRpcDiemInterface {
        let json_rpc_endpoint = format!("http://127.0.0.1:{}", swarm.get_client_port(node_index));
        JsonRpcDiemInterface::new(json_rpc_endpoint)
    }

    /// Returns a Diem Debugger pointing to a node at the given node index.
    pub fn get_diem_debugger(swarm: &DiemSwarm, node_index: usize) -> DiemDebugger {
        let (node_config, _) = load_node_config(swarm, node_index);
        let swarm_rpc_endpoint =
            format!("http://localhost:{}", node_config.json_rpc.address.port());
        DiemDebugger::json_rpc(swarm_rpc_endpoint.as_str()).unwrap()
    }

    /// Returns a Diem Event Fetcher pointing to a node at the given node index.
    pub fn get_diem_event_fetcher(swarm: &DiemSwarm, node_index: usize) -> DiemEventsFetcher {
        let (node_config, _) = load_node_config(swarm, node_index);
        let swarm_rpc_endpoint =
            format!("http://localhost:{}", node_config.json_rpc.address.port());
        DiemEventsFetcher::new(swarm_rpc_endpoint.as_str()).unwrap()
    }

    /// Returns an operational tool pointing to a validator node at the given node index.
    pub fn get_op_tool(swarm: &DiemSwarm, node_index: usize) -> OperationalTool {
        OperationalTool::new(
            format!("http://127.0.0.1:{}", swarm.get_client_port(node_index)),
            ChainId::test(),
        )
    }

    /// Loads the nodes's storage backend identified by the node index in the given swarm.
    pub fn load_backend_storage(swarm: &DiemSwarm, node_index: usize) -> SecureBackend {
        let (node_config, _) = load_node_config(swarm, node_index);
        fetch_backend_storage(&node_config, None)
    }

    /// Loads the diem root's storage backend identified by the node index in the given swarm.
    pub fn load_diem_root_storage(swarm: &DiemSwarm, node_index: usize) -> SecureBackend {
        let (node_config, _) = load_node_config(swarm, node_index);
        fetch_backend_storage(&node_config, Some("diem_root".to_string()))
    }

    /// Loads the node config for the validator at the specified index. Also returns the node
    /// config path.
    pub fn load_node_config(swarm: &DiemSwarm, node_index: usize) -> (NodeConfig, PathBuf) {
        let node_config_path = swarm.config.config_files.get(node_index).unwrap();
        let node_config = NodeConfig::load(&node_config_path).unwrap();
        (node_config, node_config_path.clone())
    }

    /// Saves the node config for the node at the specified index in the given swarm.
    pub fn save_node_config(node_config: &mut NodeConfig, swarm: &DiemSwarm, node_index: usize) {
        let node_config_path = swarm.config.config_files.get(node_index).unwrap();
        node_config.save(node_config_path).unwrap();
    }

    pub fn insert_waypoint(node_config: &mut NodeConfig, waypoint: Waypoint) {
        let f = |backend: &SecureBackend| {
            let mut storage: Storage = backend.into();<commit_msg>smoke-test: use swarm's root storage<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::smoke_test_environment::SmokeTestEnvironment;
use cli::client_proxy::ClientProxy;
use diem_config::config::{Identity, NodeConfig, SecureBackend};
use diem_crypto::ed25519::Ed25519PublicKey;
use diem_types::account_address::AccountAddress;
use rust_decimal::{prelude::FromPrimitive, Decimal};
use std::{collections::BTreeMap, fs::File, io::Write, path::PathBuf, str::FromStr};

// TODO(joshlind): Refactor all of these so that they can be contained within the calling
// test files and not shared across all tests.
pub fn compare_balances(
    expected_balances: Vec<(f64, String)>,
    extracted_balances: Vec<String>,
) -> bool {
    if extracted_balances.len() != extracted_balances.len() {
        return false;
    }

    let extracted_balances_dec: BTreeMap<_, _> = extracted_balances
        .into_iter()
        .map(|balance_str| {
            let (currency_code, stripped_str) = if balance_str.ends_with("XUS") {
                ("XUS", balance_str.trim_end_matches("XUS"))
            } else if balance_str.ends_with("XDX") {
                ("XDX", balance_str.trim_end_matches("XDX"))
            } else {
                panic!("Unexpected currency type returned for balance")
            };
            (currency_code, Decimal::from_str(stripped_str).ok())
        })
        .collect();

    expected_balances
        .into_iter()
        .all(|(balance, currency_code)| {
            if let Some(extracted_balance) = extracted_balances_dec.get(currency_code.as_str()) {
                Decimal::from_f64(balance) == *extracted_balance
            } else {
                false
            }
        })
}

/// Sets up a SmokeTestEnvironment with specified size and connects a client
/// proxy to the node_index.
pub fn setup_swarm_and_client_proxy(
    num_nodes: usize,
    node_index: usize,
) -> (SmokeTestEnvironment, ClientProxy) {
    let mut env = SmokeTestEnvironment::new(num_nodes);
    env.validator_swarm.launch();

    let client = env.get_validator_client(node_index, None);
    (env, client)
}

/// Waits for a transaction to be processed by all validator nodes in the smoke
/// test environment.
pub fn wait_for_transaction_on_all_nodes(
    env: &SmokeTestEnvironment,
    num_nodes: usize,
    account: AccountAddress,
    sequence_number: u64,
) {
    for i in 0..num_nodes {
        let client = env.get_validator_client(i, None);
        client
            .wait_for_transaction(account, sequence_number)
            .unwrap();
    }
}

/// This module provides useful functions for operating, handling and managing
/// DiemSwarm instances. It is particularly useful for working with tests that
/// require a SmokeTestEnvironment, as it provides a generic interface across
/// DiemSwarms, regardless of if the swarm is a validator swarm, validator full
/// node swarm, or a public full node swarm.
pub mod diem_swarm_utils {
    use crate::test_utils::fetch_backend_storage;
    use cli::client_proxy::ClientProxy;
    use diem_config::config::{NodeConfig, SecureBackend, WaypointConfig};
    use diem_events_fetcher::DiemEventsFetcher;
    use diem_key_manager::diem_interface::JsonRpcDiemInterface;
    use diem_operational_tool::test_helper::OperationalTool;
    use diem_secure_storage::{KVStorage, Storage};
    use diem_swarm::swarm::DiemSwarm;
    use diem_transaction_replay::DiemDebugger;
    use diem_types::{chain_id::ChainId, waypoint::Waypoint};
    use std::path::PathBuf;

    /// Returns a new client proxy connected to the given swarm at the specified
    /// node index.
    pub fn get_client_proxy(
        swarm: &DiemSwarm,
        node_index: usize,
        diem_root_key_path: &str,
        mnemonic_file_path: PathBuf,
        waypoint: Option<Waypoint>,
    ) -> ClientProxy {
        let port = swarm.get_client_port(node_index);

        let mnemonic_file_path = mnemonic_file_path
            .canonicalize()
            .expect("Unable to get canonical path of mnemonic_file_path")
            .to_str()
            .unwrap()
            .to_string();

        ClientProxy::new(
            ChainId::test(),
            &format!("http://localhost:{}/v1", port),
            &diem_root_key_path,
            &diem_root_key_path,
            &diem_root_key_path,
            false,
            /* faucet server */ None,
            Some(mnemonic_file_path),
            waypoint.unwrap_or(swarm.config.waypoint),
            true,
        )
        .unwrap()
    }

    /// Returns the JSON RPC url pointing to a node at the given
    /// node index.
    pub fn get_json_rpc_url(swarm: &DiemSwarm, node_index: usize) -> String {
        format!("http://127.0.0.1:{}", swarm.get_client_port(node_index))
    }

    /// Returns a JSON RPC based Diem Interface pointing to a node at the given
    /// node index.
    pub fn get_json_rpc_diem_interface(
        swarm: &DiemSwarm,
        node_index: usize,
    ) -> JsonRpcDiemInterface {
        let json_rpc_endpoint = format!("http://127.0.0.1:{}", swarm.get_client_port(node_index));
        JsonRpcDiemInterface::new(json_rpc_endpoint)
    }

    /// Returns a Diem Debugger pointing to a node at the given node index.
    pub fn get_diem_debugger(swarm: &DiemSwarm, node_index: usize) -> DiemDebugger {
        let (node_config, _) = load_node_config(swarm, node_index);
        let swarm_rpc_endpoint =
            format!("http://localhost:{}", node_config.json_rpc.address.port());
        DiemDebugger::json_rpc(swarm_rpc_endpoint.as_str()).unwrap()
    }

    /// Returns a Diem Event Fetcher pointing to a node at the given node index.
    pub fn get_diem_event_fetcher(swarm: &DiemSwarm, node_index: usize) -> DiemEventsFetcher {
        let (node_config, _) = load_node_config(swarm, node_index);
        let swarm_rpc_endpoint =
            format!("http://localhost:{}", node_config.json_rpc.address.port());
        DiemEventsFetcher::new(swarm_rpc_endpoint.as_str()).unwrap()
    }

    /// Returns an operational tool pointing to a validator node at the given node index.
    pub fn get_op_tool(swarm: &DiemSwarm, node_index: usize) -> OperationalTool {
        OperationalTool::new(
            format!("http://127.0.0.1:{}", swarm.get_client_port(node_index)),
            ChainId::test(),
        )
    }

    /// Loads the nodes's storage backend identified by the node index in the given swarm.
    pub fn load_backend_storage(swarm: &DiemSwarm, node_index: usize) -> SecureBackend {
        let (node_config, _) = load_node_config(swarm, node_index);
        fetch_backend_storage(&node_config, None)
    }

    /// Loads the diem root's storage backend identified by the node index in the given swarm.
    pub fn load_diem_root_storage(swarm: &DiemSwarm, _node_index: usize) -> SecureBackend {
        SecureBackend::OnDiskStorage(swarm.config.root_storage.clone())
    }

    /// Loads the node config for the validator at the specified index. Also returns the node
    /// config path.
    pub fn load_node_config(swarm: &DiemSwarm, node_index: usize) -> (NodeConfig, PathBuf) {
        let node_config_path = swarm.config.config_files.get(node_index).unwrap();
        let node_config = NodeConfig::load(&node_config_path).unwrap();
        (node_config, node_config_path.clone())
    }

    /// Saves the node config for the node at the specified index in the given swarm.
    pub fn save_node_config(node_config: &mut NodeConfig, swarm: &DiemSwarm, node_index: usize) {
        let node_config_path = swarm.config.config_files.get(node_index).unwrap();
        node_config.save(node_config_path).unwrap();
    }

    pub fn insert_waypoint(node_config: &mut NodeConfig, waypoint: Waypoint) {
        let f = |backend: &SecureBackend| {
            let mut storage: Storage = backend.into();
{'fraction_non_alphanumeric': 0.10790456300234162, 'fraction_numerical': 0.004809822163154231, 'mean_word_length': 3.021888521252227, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 6, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '961016', 'n_tokens_mistral': 4913, 'n_tokens_neox': 4867, 'n_words': 1176}