diff --git a/mithril-client/src/client.rs b/mithril-client/src/client.rs index dc04ae01c78..f1c5164fb9b 100644 --- a/mithril-client/src/client.rs +++ b/mithril-client/src/client.rs @@ -5,13 +5,15 @@ use crate::dependencies::DependenciesBuilder; use anyhow::Context; use config::builder::DefaultState; use config::ConfigBuilder; +use mithril_common::messages::MithrilStakeDistributionListMessage; use mithril_common::StdResult; use crate::common::*; -use crate::services::SnapshotService; +use crate::services::{MithrilStakeDistributionService, SnapshotService}; /// Client structure that instanciates required dependencies pub struct Client { snapshot_service: Arc, + mithril_stake_distribution_service: Arc, } impl Client { @@ -20,13 +22,20 @@ impl Client { let config_builder: ConfigBuilder = ConfigBuilder::default() .set_default("genesis_verification_key", genesis_verification_key)? .set_default("aggregator_endpoint", aggregator_endpoint)?; - let config = Arc::new(config_builder.build()?); - let snapshot_service = DependenciesBuilder::new(config) + let mut dependencies_builder = DependenciesBuilder::new(Arc::new(config_builder.build()?)); + let snapshot_service = dependencies_builder .get_snapshot_service() .await - .with_context(|| "Dependencies Builder can not get Snapshot Service")?; + .with_context(|| "Cannot instanciate Client.")?; + let mithril_stake_distribution_service = dependencies_builder + .get_mithril_stake_distribution_service() + .await + .with_context(|| "Cannot instanciate Client")?; - Ok(Self { snapshot_service }) + Ok(Self { + snapshot_service, + mithril_stake_distribution_service, + }) } /// Call the snapshot service to get a snapshot message from a digest @@ -35,16 +44,28 @@ impl Client { } /// Call the snapshot service to get the list of available snapshots - pub async fn list_snapshot(&self) -> StdResult> { + pub async fn list_snapshots(&self) -> StdResult> { self.snapshot_service.list().await } + + /// Call the mithril stake distribution service for the list of available mithril stake distributions + pub async fn list_mithril_stake_distributions( + &self, + ) -> StdResult { + self.mithril_stake_distribution_service.list().await + } } #[cfg(test)] mod tests { + use std::env::temp_dir; + use warp::{self, Filter}; - use mithril_common::test_utils::test_http_server::test_http_server; + use mithril_common::{ + messages::{MithrilStakeDistributionListItemMessage, MithrilStakeDistributionMessage}, + test_utils::test_http_server::test_http_server, + }; use super::*; @@ -65,7 +86,7 @@ mod tests { } #[tokio::test] - async fn list_snapshot() -> StdResult<()> { + async fn list_snapshots() -> StdResult<()> { let snapshot_list_message = vec![ SnapshotListItemMessage::dummy(), SnapshotListItemMessage { @@ -79,10 +100,50 @@ mod tests { warp::path!("artifact" / "snapshots").map(move || json_snapshot_list_message.clone()), ); let client = Client::new("WRITE THE VKEY HERE", &http_server.url()).await?; - let response = client.list_snapshot().await?; + let response = client.list_snapshots().await?; assert_eq!(snapshot_list_message, response); Ok(()) } + + #[tokio::test] + async fn list_mithril_stake_ditrbution() -> StdResult<()> { + let mithril_stake_distribution_list_message: MithrilStakeDistributionListMessage = vec![ + MithrilStakeDistributionListItemMessage::dummy(), + MithrilStakeDistributionListItemMessage { + hash: "12345".to_string(), + ..MithrilStakeDistributionListItemMessage::dummy() + }, + ]; + let mithril_stake_distribution_list_message_json = + serde_json::to_string(&mithril_stake_distribution_list_message)?; + let http_server = test_http_server( + warp::path!("artifact" / "mithril-stake-distributions") + .map(move || mithril_stake_distribution_list_message_json.clone()), + ); + let client = Client::new("WRITE THE VKEY HERE", &http_server.url()).await?; + let response = client.list_mithril_stake_distributions().await?; + + assert_eq!(mithril_stake_distribution_list_message, response); + + Ok(()) + } + + async fn download_mithril_stake_distribution() -> StdResult<()> { + let mithril_stake_distribution = MithrilStakeDistributionMessage::dummy(); + let mithril_stake_distribution_message_json = + serde_json::to_string(&mithril_stake_distribution)?; + let http_server = test_http_server( + warp::path!("artifact" / "mithril-stake-distribution" / String) + .map(move |msd_hash| mithril_stake_distribution_message_json.clone()), + ); + let client = Client::new("WRITE THE VKEY HERE", &http_server.url()).await?; + let download_dir = temp_dir().join("download_mithril_stake_distribution"); + let response = client + .download_mithril_stake_distribution(&download_dir) + .await?; + + Ok(()) + } }