mithril_aggregator/services/signature_consumer/
fake.rs

1use std::{collections::VecDeque, future};
2
3use async_trait::async_trait;
4use tokio::sync::Mutex;
5
6use mithril_common::{
7    entities::{SignedEntityType, SingleSignature},
8    StdResult,
9};
10
11use super::SignatureConsumer;
12
13type SignatureBatchResult = StdResult<Vec<(SingleSignature, SignedEntityType)>>;
14
15/// A fake implementation of the [SignatureConsumer] trait (test only).
16pub struct FakeSignatureConsumer {
17    signature_batches: Mutex<VecDeque<SignatureBatchResult>>,
18}
19
20impl FakeSignatureConsumer {
21    /// Creates a new `FakeSignatureConsumer` instance.
22    pub fn new(signature_batches: Vec<SignatureBatchResult>) -> Self {
23        Self {
24            signature_batches: Mutex::new(signature_batches.into()),
25        }
26    }
27}
28
29#[async_trait]
30impl SignatureConsumer for FakeSignatureConsumer {
31    async fn get_signatures(&self) -> SignatureBatchResult {
32        let mut signature_batches = self.signature_batches.lock().await;
33        match signature_batches.pop_front() {
34            None => future::pending().await,
35            Some(signature_batch) => {
36                return signature_batch;
37            }
38        }
39    }
40    fn get_origin_tag(&self) -> String {
41        "FAKE".to_string()
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    use mithril_common::{entities::Epoch, test_utils::fake_data};
50
51    #[tokio::test]
52    async fn fake_signature_consumer_returns_signature_batches_in_expected_order() {
53        let consumer = FakeSignatureConsumer::new(vec![
54            Ok(vec![(
55                fake_data::single_signature(vec![1, 2, 3]),
56                SignedEntityType::MithrilStakeDistribution(Epoch(1)),
57            )]),
58            Ok(vec![(
59                fake_data::single_signature(vec![4, 5, 6]),
60                SignedEntityType::MithrilStakeDistribution(Epoch(2)),
61            )]),
62        ]);
63
64        let result = consumer.get_signatures().await.unwrap();
65        assert_eq!(
66            vec![(
67                fake_data::single_signature(vec![1, 2, 3]),
68                SignedEntityType::MithrilStakeDistribution(Epoch(1)),
69            )],
70            result
71        );
72
73        let result = consumer.get_signatures().await.unwrap();
74        assert_eq!(
75            vec![(
76                fake_data::single_signature(vec![4, 5, 6]),
77                SignedEntityType::MithrilStakeDistribution(Epoch(2)),
78            )],
79            result
80        );
81    }
82}