Add a busy waiting of signed entity type for integration tests

input-output-hk · Jul 15, 2024 · a21746c · a21746c
1 parent f56a690
commit a21746c
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 0 deletions.
diff --git a/mithril-aggregator/tests/test_extensions/aggregator_observer.rs b/mithril-aggregator/tests/test_extensions/aggregator_observer.rs
@@ -112,4 +112,34 @@ impl AggregatorObserver {
             .signed_entity_config
             .time_point_to_signed_entity(discriminant, &time_point))
     }
+
+    pub async fn is_last_signed_entity(
+        &self,
+        signed_entity_type_expected: &SignedEntityType,
+    ) -> StdResult<bool> {
+        match signed_entity_type_expected {
+            SignedEntityType::CardanoImmutableFilesFull(_) => Ok(Some(signed_entity_type_expected)
+                == self
+                    .signed_entity_service
+                    .get_last_signed_snapshots(1)
+                    .await?
+                    .first()
+                    .map(|s| &s.signed_entity_type)),
+            SignedEntityType::MithrilStakeDistribution(_) => Ok(Some(signed_entity_type_expected)
+                == self
+                    .signed_entity_service
+                    .get_last_signed_mithril_stake_distributions(1)
+                    .await?
+                    .first()
+                    .map(|s| &s.signed_entity_type)),
+            SignedEntityType::CardanoTransactions(_, _) => Ok(Some(signed_entity_type_expected)
+                == self
+                    .signed_entity_service
+                    .get_last_cardano_transaction_snapshot()
+                    .await?
+                    .map(|s| s.signed_entity_type)
+                    .as_ref()),
+            _ => Ok(false),
+        }
+    }
 }
diff --git a/mithril-aggregator/tests/test_extensions/expected_certificate.rs b/mithril-aggregator/tests/test_extensions/expected_certificate.rs
@@ -50,4 +50,8 @@ impl ExpectedCertificate {
     pub fn genesis_identifier(beacon: &CardanoDbBeacon) -> String {
         format!("genesis-{:?}", beacon)
     }
+
+    pub fn get_signed_type(&self) -> Option<SignedEntityType> {
+        self.signed_type.clone()
+    }
 }
diff --git a/mithril-aggregator/tests/test_extensions/runtime_tester.rs b/mithril-aggregator/tests/test_extensions/runtime_tester.rs
@@ -9,6 +9,7 @@ use mithril_aggregator::{
     AggregatorRuntime, Configuration, DependencyContainer, DumbSnapshotUploader, DumbSnapshotter,
     SignerRegistrationError,
 };
+use mithril_common::entities::SignedEntityType;
 use mithril_common::{
     cardano_block_scanner::{DumbBlockScanner, ScannedBlock},
     chain_observer::{ChainObserver, FakeObserver},
@@ -51,6 +52,13 @@ macro_rules! cycle_err {
 #[macro_export]
 macro_rules! assert_last_certificate_eq {
     ( $tester:expr, $expected_certificate:expr ) => {{
+        if let Some(signed_type) = $expected_certificate.get_signed_type() {
+            $tester
+                .wait_until_signed_entity(&signed_type)
+                .await
+                .unwrap();
+        }
+
         let last_certificate = RuntimeTester::get_last_expected_certificate(&mut $tester)
             .await
             .unwrap();
@@ -568,4 +576,28 @@ impl RuntimeTester {
 
         Ok(cert_identifier)
     }
+
+    /// Wait until the last stored signed entity of the given type
+    /// corresponds to the expected signed entity type
+    pub async fn wait_until_signed_entity(
+        &self,
+        signed_entity_type_expected: &SignedEntityType,
+    ) -> StdResult<()> {
+        let mut max_iteration = 100;
+        while !self
+            .observer
+            .is_last_signed_entity(signed_entity_type_expected)
+            .await?
+        {
+            max_iteration -= 1;
+            if max_iteration <= 0 {
+                return Err(anyhow!(
+                    "Signed entity not found: {signed_entity_type_expected}"
+                ));
+            }
+            tokio::time::sleep(Duration::from_millis(1)).await;
+        }
+
+        Ok(())
+    }
 }