Fix end to end test flakiness in CI

mithril-aggregator/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-aggregator"
-version = "0.4.39"
+version = "0.4.40"
 description = "A Mithril Aggregator server"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-aggregator/src/dependency_injection/builder.rs

@@ -1,7 +1,6 @@
 use anyhow::Context;
 use semver::Version;
 use slog::Logger;
-use slog_scope::debug;
 use sqlite::Connection;
 use std::sync::Arc;
 use tokio::{
@@ -536,10 +535,34 @@ impl DependenciesBuilder {
     async fn build_protocol_parameters_store(
         &mut self,
     ) -> Result<Arc<dyn ProtocolParametersStorer>> {
-        Ok(Arc::new(EpochSettingStore::new(
+        let protocol_parameters_store = EpochSettingStore::new(
             self.get_sqlite_connection().await?,
             self.configuration.safe_epoch_retention_limit(),
-        )))
+        );
+        let current_epoch = self
+            .get_chain_observer()
+            .await?
+            .get_current_epoch()
+            .await
+            .map_err(|e| DependenciesBuilderError::Initialization {
+                message: "cannot create aggregator runner: failed to retrieve current epoch."
+                    .to_string(),
+                error: Some(e.into()),
+            })?
+            .ok_or(DependenciesBuilderError::Initialization {
+                message: "cannot build aggregator runner: no epoch returned.".to_string(),
+                error: None,
+            })?;
+
+        protocol_parameters_store
+            .handle_discrepancies_at_startup(current_epoch, &self.configuration.protocol_parameters)
+            .await
+            .map_err(|e| DependenciesBuilderError::Initialization {
+                message: "can not create aggregator runner".to_string(),
+                error: Some(e),
+            })?;
+
+        Ok(Arc::new(protocol_parameters_store))
     }
 
     /// Get a configured [ProtocolParametersStorer].
@@ -1214,57 +1237,6 @@ impl DependenciesBuilder {
 
     /// Create the AggregatorRunner
     pub async fn create_aggregator_runner(&mut self) -> Result<AggregatorRuntime> {
-        // initialize ProtocolParameters store if needed
-        {
-            let current_epoch = self
-                .get_chain_observer()
-                .await?
-                .get_current_epoch()
-                .await
-                .map_err(|e| DependenciesBuilderError::Initialization {
-                    message: "cannot create aggregator runner: failed to retrieve current epoch."
-                        .to_string(),
-                    error: Some(e.into()),
-                })?
-                .ok_or(DependenciesBuilderError::Initialization {
-                    message: "cannot build aggregator runner: no epoch returned.".to_string(),
-                    error: None,
-                })?;
-            let (work_epoch, epoch_to_sign, next_epoch_to_sign) = match current_epoch {
-                Epoch(0) => (Epoch(0), Epoch(1), Epoch(2)),
-                epoch => (
-                    epoch.offset_to_signer_retrieval_epoch().unwrap(),
-                    epoch.offset_to_next_signer_retrieval_epoch(),
-                    epoch.offset_to_next_signer_retrieval_epoch().next(),
-                ),
-            };
-            let protocol_parameters_store = self.get_protocol_parameters_store().await?;
-
-            if protocol_parameters_store
-                .get_protocol_parameters(work_epoch)
-                .await
-                .map_err(|e| DependenciesBuilderError::Initialization {
-                    message: "can not create aggregator runner".to_string(),
-                    error: Some(e),
-                })?
-                .is_none()
-            {
-                debug!("First launch, will use the configured protocol parameters for the current and next epoch certificate");
-
-                for epoch in [work_epoch, epoch_to_sign, next_epoch_to_sign] {
-                    protocol_parameters_store
-                        .save_protocol_parameters(
-                            epoch,
-                            self.configuration.protocol_parameters.clone(),
-                        )
-                        .await
-                        .map_err(|e| DependenciesBuilderError::Initialization {
-                            message: "can not create aggregator runner".to_string(),
-                            error: Some(e),
-                        })?;
-                }
-            }
-        }
         let dependency_container = Arc::new(self.build_dependency_container().await?);
 
         let config = AggregatorConfig::new(

mithril-aggregator/src/lib.rs

@@ -55,8 +55,7 @@ pub use snapshotter::{
     SnapshotterCompressionAlgorithm,
 };
 pub use store::{
-    CertificatePendingStore, ProtocolParametersStore, ProtocolParametersStorer,
-    VerificationKeyStore, VerificationKeyStorer,
+    CertificatePendingStore, ProtocolParametersStorer, VerificationKeyStore, VerificationKeyStorer,
 };
 pub use tools::{
     CExplorerSignerRetriever, SignersImporter, SignersImporterPersister, SignersImporterRetriever,

mithril-aggregator/src/services/epoch_service.rs

@@ -467,7 +467,8 @@ mod tests {
     use std::collections::{BTreeSet, HashMap};
 
     use crate::services::epoch_service::tests::ServiceBuilderParameters::WithFutureProtocolParameters;
-    use crate::{ProtocolParametersStore, VerificationKeyStore};
+    use crate::store::FakeProtocolParametersStorer;
+    use crate::VerificationKeyStore;
 
     use super::*;
 
@@ -558,26 +559,20 @@ mod tests {
             }
         }
 
-        let protocol_parameters_store = ProtocolParametersStore::new(
-            Box::new(
-                MemoryAdapter::new(Some(vec![
-                    (
-                        signer_retrieval_epoch,
-                        current_epoch_fixture.protocol_parameters(),
-                    ),
-                    (
-                        next_signer_retrieval_epoch,
-                        next_epoch_fixture.protocol_parameters(),
-                    ),
-                    (
-                        next_signer_retrieval_epoch.next(),
-                        upcoming_protocol_parameters.clone(),
-                    ),
-                ]))
-                .unwrap(),
+        let protocol_parameters_store = FakeProtocolParametersStorer::new(vec![
+            (
+                signer_retrieval_epoch,
+                current_epoch_fixture.protocol_parameters(),
             ),
-            None,
-        );
+            (
+                next_signer_retrieval_epoch,
+                next_epoch_fixture.protocol_parameters(),
+            ),
+            (
+                next_signer_retrieval_epoch.next(),
+                upcoming_protocol_parameters.clone(),
+            ),
+        ]);
         let vkey_store = VerificationKeyStore::new(Box::new(
             MemoryAdapter::new(Some(vec![
                 (

mithril-aggregator/src/store/mod.rs

@@ -3,9 +3,11 @@ mod protocol_parameters_store;
 mod verification_key_store;
 
 pub use pending_certificate_store::CertificatePendingStore;
-pub use protocol_parameters_store::{ProtocolParametersStore, ProtocolParametersStorer};
+pub use protocol_parameters_store::ProtocolParametersStorer;
 pub use verification_key_store::{VerificationKeyStore, VerificationKeyStorer};
 
+#[cfg(test)]
+pub use protocol_parameters_store::FakeProtocolParametersStorer;
 #[cfg(test)]
 pub use verification_key_store::test_suite as verification_key_store_test_suite;
 #[cfg(test)]