Fix flaky end to end tests

mithril-aggregator/Cargo.toml

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

mithril-aggregator/src/runtime/error.rs

@@ -60,11 +60,15 @@ impl From<Box<dyn StdError + Sync + Send>> for RuntimeError {
 // TODO: Are these errors still relevant, do we need to remove them?
 #[allow(clippy::enum_variant_names)]
 pub enum RunnerError {
-    /// No stack distribution found
-    #[error("Missing stack distribution: '{0}'.")]
+    /// No stake distribution found
+    #[error("Missing stake distribution: '{0}'.")]
     MissingStakeDistribution(String),
 
     /// Missing protocol parameters
     #[error("Missing protocol parameters: '{0}'.")]
     MissingProtocolParameters(String),
+
+    /// Missing next aggregate verification key
+    #[error("Missing next aggregate verification key: '{0}'.")]
+    MissingNextAggregateVerificationKey(String),
 }

mithril-aggregator/src/runtime/runner.rs

@@ -211,6 +211,11 @@ impl AggregatorRunnerTrait for AggregatorRunner {
                     "RUNNER: get_current_non_certified_open_message_for_signed_entity_type: no open message found, a new one will be created";
                     "signed_entity_type" => ?signed_entity_type
                 );
+                // In order to fix flakiness in the end to end tests, we are compelled to update the beacon (of the multi-signer)
+                // This avoids the computation of the next AVK on the wrong epoch (in 'compute_protocol_message')
+                // TODO: remove 'current_beacon' from the multi-signer state which will avoid this fix
+                let chain_beacon: Beacon = self.get_beacon_from_chain().await?;
+                self.update_beacon(&chain_beacon).await?;
                 let protocol_message = self.compute_protocol_message(signed_entity_type).await?;
                 Some(
                     self.create_open_message(signed_entity_type, &protocol_message)
@@ -385,6 +390,7 @@ impl AggregatorRunnerTrait for AggregatorRunner {
         debug!("RUNNER: update protocol parameters"; "beacon" => #?new_beacon);
         let protocol_parameters = self.dependencies.config.protocol_parameters.clone();
 
+        self.update_beacon(new_beacon).await?;
         self.dependencies
             .multi_signer
             .write()
@@ -411,7 +417,11 @@ impl AggregatorRunnerTrait for AggregatorRunner {
             multi_signer
                 .compute_next_stake_distribution_aggregate_verification_key()
                 .await?
-                .unwrap_or_default(),
+                .ok_or_else(|| {
+                    RunnerError::MissingNextAggregateVerificationKey(format!(
+                        "Next aggregate verification key could not be computed for signer endity type {signed_entity_type:?}"
+                    ))
+                })?,
         );
 
         Ok(protocol_message)

mithril-signer/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-signer"
-version = "0.2.54"
+version = "0.2.55"
 description = "A Mithril Signer"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-signer/src/runtime/state_machine.rs

@@ -237,7 +237,7 @@ impl StateMachine {
                         info!(" ⋅ pending certificate has not changed, waiting…");
                     } else {
                         info!(" → new pending certificate detected, transiting to REGISTERED");
-                        self.state = self.transition_from_signed_to_registered().await?;
+                        self.state = self.transition_from_signed_to_registered(*epoch).await?;
                     }
                 } else {
                     info!(" ⋅ no pending certificate, waiting…");
@@ -291,14 +291,11 @@ impl StateMachine {
         Ok(SignerState::Unregistered { epoch })
     }
 
-    async fn transition_from_signed_to_registered(&self) -> Result<SignerState, RuntimeError> {
-        let current_beacon = self
-            .get_current_beacon("checking if epoch has changed")
-            .await?;
-
-        Ok(SignerState::Registered {
-            epoch: current_beacon.epoch,
-        })
+    async fn transition_from_signed_to_registered(
+        &self,
+        epoch: Epoch,
+    ) -> Result<SignerState, RuntimeError> {
+        Ok(SignerState::Registered { epoch })
     }
 
     async fn transition_from_registered_to_unregistered(
@@ -691,7 +688,7 @@ mod tests {
         let mut runner = MockSignerRunner::new();
         runner
             .expect_get_current_beacon()
-            .times(2)
+            .times(1)
             .returning(move || Ok(beacon_clone.to_owned()));
         runner
             .expect_get_pending_certificate()

mithril-test-lab/mithril-end-to-end/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-end-to-end"
-version = "0.1.24"
+version = "0.1.25"
 authors = { workspace = true }
 edition = { workspace = true }
 documentation = { workspace = true }

mithril-test-lab/mithril-end-to-end/src/end_to_end_spec.rs

@@ -259,8 +259,8 @@ async fn update_protocol_parameters(aggregator: &mut Aggregator) -> Result<(), S
     aggregator.stop().await?;
     let protocol_parameters_new = ProtocolParameters {
         k: 150,
-        m: 200,
-        phi_f: 0.85,
+        m: 210,
+        phi_f: 0.80,
     };
     info!(
         "> updating protocol parameters to {:?}...",