diff --git a/test_fixtures/src/lib.rs b/test_fixtures/src/lib.rs index 390c7b7f3..e112b93b7 100644 --- a/test_fixtures/src/lib.rs +++ b/test_fixtures/src/lib.rs @@ -6,7 +6,7 @@ use std::{collections::HashMap, net::SocketAddr, path::Path, sync::LazyLock}; use anyhow::{Context as _, Result}; use common::{HashType, transaction::LeeTransaction}; use futures::FutureExt as _; -use indexer_service::IndexerHandle; +use indexer_service::{ChannelId, IndexerHandle}; use lee::{AccountId, PrivacyPreservingTransaction}; use lee_core::Commitment; use log::{debug, error}; @@ -71,73 +71,149 @@ pub struct DiskSizes { pub wallet_bytes: u64, } +pub struct SequencerComponent { + /// In fact, not optional, just for Drop implementation. + sequencer_handle: Option, + temp_sequencer_dir: TempDir, + sequencer_client: SequencerClient, +} + +pub struct TestContextZone { + /// Every zone must have its own wallet, otherwise multi-sequecner client loses sence. + wallet: WalletCore, + wallet_password: String, + temp_wallet_dir: TempDir, + config: MultiNodeTestContextConfig, + /// Each sequencer component is mapped from sequencer `SocketAddr`. + sequencer_components: HashMap, + indexer_components: Option, +} + /// Test context which sets up a sequencer and a wallet for integration tests. /// /// It's memory and logically safe to create multiple instances of this struct in parallel tests, /// as each instance uses its own temporary directories for sequencer and wallet data. // NOTE: Order of fields is important for proper drop order. pub struct TestContext { - sequencer_clients: HashMap, - wallet: WalletCore, - wallet_password: String, - /// Optional to move out value in Drop. - sequencer_handles: Option>, - indexer_components: Option, + zones: HashMap, bedrock_compose: DockerCompose, bedrock_addr: SocketAddr, - temp_sequencer_dirs: Vec, - temp_wallet_dir: TempDir, - config: MultiNodeTestContextConfig, } impl TestContext { /// Create new test context. - pub async fn new_custom(config: MultiNodeTestContextConfig) -> Result { - Self::builder(config).build().await + pub async fn new_custom(configs: Vec) -> Result { + Self::builder(configs).build().await } - /// Create new test context with default config. + /// Create new test context with default config(1 zone). pub async fn new() -> Result { - Self::builder(MultiNodeTestContextConfig::default()) + Self::builder(vec![MultiNodeTestContextConfig::default()]) .build() .await } /// Get a builder for the test context to customize its configuration. #[must_use] - pub fn builder(config: MultiNodeTestContextConfig) -> TestContextBuilder { - TestContextBuilder::new(config) + pub fn builder(configs: Vec) -> MultiZoneTestContextBuilder { + MultiZoneTestContextBuilder { + zone_builders: configs.into_iter().map(TestContextBuilder::new).collect(), + } } - /// Get reference to the wallet. - #[must_use] - pub const fn wallet(&self) -> &WalletCore { - &self.wallet - } - - #[must_use] - pub fn wallet_password(&self) -> &str { - &self.wallet_password - } - - /// Get mutable reference to the wallet. - pub const fn wallet_mut(&mut self) -> &mut WalletCore { - &mut self.wallet - } - - /// Get reference to the sequencer client in default case (1 sequencer). - #[must_use] - pub fn sequencer_client(&self) -> &SequencerClient { - self.sequencer_clients + /// Reference for the default zone(in case if only one present). + fn default_zone(&self) -> &TestContextZone { + self.zones .values() .next() - .expect("Must be at least one sequencer client") + .expect("Must be at least one zone") + } + + /// Reference for the default sequencer component(in case, if only one zone exists and only + /// one sequencer exists). + fn default_sequencer_component(&self) -> &SequencerComponent { + self.default_zone() + .sequencer_components + .values() + .next() + .expect("Must be at least one integration component") + } + + /// Mutable reference for the default zone(in case if only one present). + fn default_zone_mut(&mut self) -> &mut TestContextZone { + self.zones + .values_mut() + .next() + .expect("Must be at least one zone") + } + + /// Mutable reference for the default sequencer component(in case, if only one zone exists and + /// only one sequencer exists). + fn default_sequencer_component_mut(&mut self) -> &mut SequencerComponent { + self.default_zone_mut() + .sequencer_components + .values_mut() + .next() + .expect("Must be at least one integration component") + } + + /// Get reference to the deafault wallet. + #[must_use] + pub fn wallet(&self) -> &WalletCore { + &self.default_zone().wallet + } + + /// Get password of the default wallet. + #[must_use] + pub fn wallet_password(&self) -> &str { + &self.default_zone().wallet_password + } + + /// Get mutable reference to default the wallet. + pub fn wallet_mut(&mut self) -> &mut WalletCore { + &mut self.default_zone_mut().wallet + } + + /// Get reference to the zone wallet. + #[must_use] + pub fn wallet_zone(&self, channel_id: ChannelId) -> Option<&WalletCore> { + self.zones.get(&channel_id).map(|val| &val.wallet) + } + + /// Get password of the zone wallet. + #[must_use] + pub fn wallet_password_zone(&self, channel_id: ChannelId) -> Option<&str> { + self.zones + .get(&channel_id) + .map(|val| val.wallet_password.as_str()) + } + + /// Get mutable reference to the zone wallet. + pub fn wallet_mut_zone(&mut self, channel_id: ChannelId) -> Option<&mut WalletCore> { + self.zones.get_mut(&channel_id).map(|val| &mut val.wallet) + } + + /// Get reference to the sequencer client in default case (1 zone, 1 sequencer). + #[must_use] + pub fn sequencer_client(&self) -> &SequencerClient { + &self.default_sequencer_component().sequencer_client } /// Get reference to the sequencer client. #[must_use] - pub fn sequencer_client_getter(&self, addr: &SocketAddr) -> Option<&SequencerClient> { - self.sequencer_clients.get(addr) + pub fn sequencer_client_getter( + &self, + channel_id: ChannelId, + addr: &SocketAddr, + ) -> Option<&SequencerClient> { + self.zones + .get(&channel_id) + .map(|val| { + val.sequencer_components + .get(addr) + .map(|vall| &vall.sequencer_client) + }) + .flatten() } /// Get the Bedrock Node address. @@ -146,7 +222,7 @@ impl TestContext { self.bedrock_addr } - /// Get reference to the indexer. + /// Get reference to the default indexer(1 zone). /// /// # Panics /// @@ -154,13 +230,15 @@ impl TestContext { /// [`TestContextBuilder::disable_indexer()`]. #[must_use] pub fn indexer(&self) -> &IndexerHandle { - self.indexer_components + &self + .default_zone() + .indexer_components .as_ref() - .map(|components| &components.indexer_handle) .expect("Called `TestContext::indexer()` on context with disabled indexer") + .indexer_handle } - /// Get the indexer's bound socket address. + /// Get the default indexer's(1 zone) bound socket address. /// /// # Panics /// @@ -170,7 +248,7 @@ impl TestContext { self.indexer().addr() } - /// Get reference to the indexer client. + /// Get reference to the default indexer(1 zone) client. /// /// # Panics /// @@ -178,16 +256,64 @@ impl TestContext { /// [`TestContextBuilder::disable_indexer()`]. #[must_use] pub fn indexer_client(&self) -> &IndexerClient { - self.indexer_components + &self + .default_zone() + .indexer_components .as_ref() - .map(|components| &components.indexer_client) - .expect("Called `TestContext::indexer_client()` on context with disabled indexer") + .expect("Called `TestContext::indexer()` on context with disabled indexer") + .indexer_client + } + + /// Get reference to the indexer. + /// + /// # Panics + /// + /// Panics if the indexer is not enabled in the test context. See + /// [`TestContextBuilder::disable_indexer()`]. + #[must_use] + pub fn indexer_getter(&self, channel_id: ChannelId) -> Option<&IndexerHandle> { + self.zones + .get(&channel_id) + .map(|val| { + val.indexer_components + .as_ref() + .map(|val| &val.indexer_handle) + }) + .flatten() + } + + /// Get the default indexer's bound socket address. + /// + /// # Panics + /// + /// Panics if the indexer is not enabled in the test context. + #[must_use] + pub fn indexer_addr_getter(&self, channel_id: ChannelId) -> Option { + self.indexer_getter(channel_id).map(|val| val.addr()) + } + + /// Get reference to the indexer client. + /// + /// # Panics + /// + /// Panics if the indexer is not enabled in the test context. See + /// [`TestContextBuilder::disable_indexer()`]. + #[must_use] + pub fn indexer_client_getter(&self, channel_id: ChannelId) -> Option<&IndexerClient> { + self.zones + .get(&channel_id) + .map(|val| { + val.indexer_components + .as_ref() + .map(|val| &val.indexer_client) + }) + .flatten() } #[must_use] /// Get the multi-node config. - pub const fn config(&self) -> &MultiNodeTestContextConfig { - &self.config + pub fn config(&self, channel_id: ChannelId) -> Option<&MultiNodeTestContextConfig> { + self.zones.get(&channel_id).map(|val| &val.config) } /// Recursively-sized bytes on disk for sequencer + indexer + wallet tempdirs. @@ -195,22 +321,33 @@ impl TestContext { #[must_use] pub fn disk_sizes(&self) -> DiskSizes { DiskSizes { - sequencer_bytes: self - .temp_sequencer_dirs - .iter() - .fold(0, |acc, dir| acc.saturating_add(dir_size_bytes(dir.path()))), - indexer_bytes: self - .indexer_components - .as_ref() - .map_or(0, |c| dir_size_bytes(c.temp_dir.path())), - wallet_bytes: dir_size_bytes(self.temp_wallet_dir.path()), + sequencer_bytes: self.zones.values().fold(0, |acc, zone| { + acc.saturating_add( + zone.sequencer_components + .values() + .fold(0, |accc, component| { + accc.saturating_add(dir_size_bytes(component.temp_sequencer_dir.path())) + }), + ) + }), + indexer_bytes: self.zones.values().fold(0, |acc, zone| { + acc.saturating_add( + zone.indexer_components + .as_ref() + .map_or(0, |val| dir_size_bytes(val.temp_dir.path())), + ) + }), + wallet_bytes: self.zones.values().fold(0, |acc, zone| { + acc.saturating_add(dir_size_bytes(zone.temp_wallet_dir.path())) + }), } } - /// Get existing public account IDs in the wallet. + /// Get default(1 zone) existing public account IDs in the wallet. #[must_use] pub fn existing_public_accounts(&self) -> Vec { - self.wallet + self.default_zone() + .wallet .storage() .key_chain() .public_account_ids() @@ -218,46 +355,81 @@ impl TestContext { .collect() } - /// Get existing private account IDs in the wallet. + /// Get default(1 zone) existing private account IDs in the wallet. #[must_use] pub fn existing_private_accounts(&self) -> Vec { - self.wallet + self.default_zone() + .wallet .storage() .key_chain() .private_account_ids() .map(|(account_id, _idx)| account_id) .collect() } + + /// Get existing public account IDs in the wallet. + #[must_use] + pub fn existing_public_accounts_zone(&self, channel_id: ChannelId) -> Option> { + self.wallet_zone(channel_id).map(|wallet_ref| { + wallet_ref + .storage() + .key_chain() + .public_account_ids() + .map(|(account_id, _idx)| account_id) + .collect() + }) + } + + /// Get existing private account IDs in the wallet. + #[must_use] + pub fn existing_private_accounts_zone(&self, channel_id: ChannelId) -> Option> { + self.wallet_zone(channel_id).map(|wallet_ref| { + wallet_ref + .storage() + .key_chain() + .private_account_ids() + .map(|(account_id, _idx)| account_id) + .collect() + }) + } } impl Drop for TestContext { fn drop(&mut self) { let Self { - sequencer_handles, + zones, bedrock_compose, bedrock_addr: _, - indexer_components: _, - sequencer_clients: _, - wallet: _, - wallet_password: _, - temp_sequencer_dirs: _, - temp_wallet_dir: _, - config: _, } = self; - let sequencer_handles = sequencer_handles - .take() - .expect("Sequencer handle should be present in TestContext drop"); + for TestContextZone { + wallet: _, + wallet_password: _, + temp_wallet_dir: _, + config: _, + sequencer_components, + indexer_components: _, + } in zones.values_mut() + { + for SequencerComponent { + sequencer_handle, + temp_sequencer_dir: _, + sequencer_client: _, + } in sequencer_components.values_mut() + { + let sequencer_handle = sequencer_handle + .take() + .expect("Sequencer handle should be present in TestContext drop"); - for sequencer_handle in sequencer_handles { - if !sequencer_handle.is_healthy() { - let Err(err) = sequencer_handle - .failed() - .now_or_never() - .expect("Sequencer handle should not be running"); - error!( - "Sequencer handle has unexpectedly stopped before TestContext drop with error: {err:#}" - ); + if !sequencer_handle.is_healthy() { + let Err(err) = sequencer_handle + .failed() + .now_or_never() + .expect("Sequencer handle should not be running"); + error!( + "Sequencer handle has unexpectedly stopped before TestContext drop with error: {err:#}" + ); + } } } @@ -290,6 +462,10 @@ pub struct TestContextBuilder { config: MultiNodeTestContextConfig, } +pub struct MultiZoneTestContextBuilder { + zone_builders: Vec, +} + impl TestContextBuilder { fn new(config: MultiNodeTestContextConfig) -> Self { Self { @@ -346,7 +522,7 @@ impl TestContextBuilder { self } - pub async fn build(self) -> Result { + pub async fn build(self, bedrock_addr: SocketAddr) -> Result { let Self { genesis_transactions, sequencer_partial_config, @@ -365,10 +541,6 @@ impl TestContextBuilder { // must build live. Otherwise load the fixture (fails if it is missing). let use_prebuilt = !from_scratch && genesis_transactions.is_none(); - let (bedrock_compose, bedrock_addr) = setup_bedrock_node() - .await - .context("Failed to setup Bedrock node")?; - let indexer_components = if enable_indexer { let (indexer_handle, temp_indexer_dir) = setup_indexer(bedrock_addr, config::bedrock_channel_id(), None) @@ -481,6 +653,18 @@ impl TestContextBuilder { } } +impl MultiZoneTestContextBuilder { + pub async fn build(self) -> Result { + let (bedrock_compose, bedrock_addr) = setup_bedrock_node() + .await + .context("Failed to setup Bedrock node")?; + + self.zone_builders + .into_iter() + .map(|ctxb| TestContextBuilder::build(ctxb, bedrock_addr.clone())) + } +} + /// A test context to be used in normal #[test] tests. pub struct BlockingTestContext { ctx: Option,