mithril_aggregator/artifact_builder/
cardano_stake_distribution.rs

1use anyhow::anyhow;
2use async_trait::async_trait;
3use std::sync::Arc;
4
5use mithril_common::{
6    entities::{CardanoStakeDistribution, Certificate, Epoch},
7    signable_builder::StakeDistributionRetriever,
8    StdResult,
9};
10
11use crate::ArtifactBuilder;
12
13/// A [CardanoStakeDistributionArtifact] builder
14pub struct CardanoStakeDistributionArtifactBuilder {
15    stake_distribution_retriever: Arc<dyn StakeDistributionRetriever>,
16}
17
18impl CardanoStakeDistributionArtifactBuilder {
19    /// CardanoStakeDistribution artifact builder factory
20    pub fn new(stake_distribution_retriever: Arc<dyn StakeDistributionRetriever>) -> Self {
21        Self {
22            stake_distribution_retriever,
23        }
24    }
25}
26
27#[async_trait]
28impl ArtifactBuilder<Epoch, CardanoStakeDistribution> for CardanoStakeDistributionArtifactBuilder {
29    async fn compute_artifact(
30        &self,
31        epoch: Epoch,
32        _certificate: &Certificate,
33    ) -> StdResult<CardanoStakeDistribution> {
34        let stake_distribution = self
35            .stake_distribution_retriever
36            .retrieve(epoch.offset_to_cardano_stake_distribution_snapshot_epoch())
37            .await?
38            .ok_or_else(|| anyhow!("No stake distribution found for epoch '{}'", epoch))?;
39
40        Ok(CardanoStakeDistribution::new(epoch, stake_distribution))
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use mithril_common::{entities::StakeDistribution, test_utils::fake_data};
47    use mockall::{mock, predicate::eq};
48
49    use super::*;
50
51    mock! {
52        pub StakeDistributionRetrieverImpl {}
53
54        #[async_trait]
55        impl StakeDistributionRetriever for StakeDistributionRetrieverImpl {
56            async fn retrieve(&self, epoch: Epoch) -> StdResult<Option<StakeDistribution>>;
57        }
58    }
59
60    #[tokio::test]
61    async fn compute_artifact_returns_valid_artifact_and_retrieve_with_epoch_offset() {
62        let epoch = Epoch(1);
63        let epoch_to_retrieve = Epoch(3);
64        let certificate = fake_data::certificate("whatever".to_string());
65        let stake_distribution = StakeDistribution::from([("pool-123".to_string(), 123)]);
66        let stake_distribution_clone = stake_distribution.clone();
67        let mut mock_storer = MockStakeDistributionRetrieverImpl::new();
68        mock_storer
69            .expect_retrieve()
70            .with(eq(epoch_to_retrieve))
71            .return_once(move |_| Ok(Some(stake_distribution_clone)));
72        let builder = CardanoStakeDistributionArtifactBuilder::new(Arc::new(mock_storer));
73
74        let cardano_stake_distribution =
75            builder.compute_artifact(epoch, &certificate).await.unwrap();
76
77        let expected = CardanoStakeDistribution::new(epoch, stake_distribution);
78        assert_eq!(cardano_stake_distribution, expected);
79    }
80
81    #[tokio::test]
82    async fn compute_artifact_returns_error_if_no_stakes_found_for_epoch() {
83        let epoch = Epoch(1);
84        let epoch_to_retrieve = Epoch(3);
85        let certificate = fake_data::certificate("whatever".to_string());
86        let mut mock_storer = MockStakeDistributionRetrieverImpl::new();
87        mock_storer
88            .expect_retrieve()
89            .with(eq(epoch_to_retrieve))
90            .return_once(move |_| Ok(None));
91        let builder = CardanoStakeDistributionArtifactBuilder::new(Arc::new(mock_storer));
92
93        builder
94            .compute_artifact(epoch, &certificate)
95            .await
96            .expect_err("Should return error");
97    }
98}