mithril_signer/runtime/
state_machine.rs

1use anyhow::Error;
2use chrono::Local;
3use slog::{debug, info, Logger};
4use std::{fmt::Display, ops::Deref, sync::Arc, time::Duration};
5use tokio::sync::Mutex;
6
7use mithril_common::{
8    crypto_helper::ProtocolInitializerError,
9    entities::{Epoch, TimePoint},
10    logging::LoggerExtensions,
11};
12
13use crate::{
14    entities::{BeaconToSign, SignerEpochSettings},
15    services::AggregatorClientError,
16    MetricsService,
17};
18
19use super::{Runner, RuntimeError};
20
21/// Different possible states of the state machine.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub enum SignerState {
24    /// Starting state
25    Init,
26    /// Hold the latest known epoch in order to help synchronisation
27    /// with the aggregator
28    Unregistered {
29        /// Current Epoch
30        epoch: Epoch,
31    },
32
33    /// `ReadyToSign` state. The signer is registered and ready to sign new messages.
34    ReadyToSign {
35        /// Epoch when signer transitioned to the state.
36        epoch: Epoch,
37    },
38
39    /// `RegisteredNotAbleToSign` state. The signer is registered but not able to sign for the duration of the epoch.
40    RegisteredNotAbleToSign {
41        /// Epoch when signer transitioned to the state.
42        epoch: Epoch,
43    },
44}
45
46impl SignerState {
47    /// Returns `true` if the state in `Init`
48    pub fn is_init(&self) -> bool {
49        matches!(*self, SignerState::Init)
50    }
51
52    /// Returns `true` if the state in `Unregistered`
53    pub fn is_unregistered(&self) -> bool {
54        matches!(*self, SignerState::Unregistered { .. })
55    }
56
57    /// Returns `true` if the state in `ReadyToSign`
58    pub fn is_ready_to_sign(&self) -> bool {
59        matches!(*self, SignerState::ReadyToSign { .. })
60    }
61
62    /// Returns `true` if the state in `RegisteredNotAbleToSign`
63    pub fn is_registered_not_able_to_sign(&self) -> bool {
64        matches!(*self, SignerState::RegisteredNotAbleToSign { .. })
65    }
66}
67
68impl Display for SignerState {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        match self {
71            Self::Init => write!(f, "Init"),
72            Self::Unregistered { epoch } => write!(f, "Unregistered - {epoch:?}"),
73            Self::RegisteredNotAbleToSign { epoch } => {
74                write!(f, "RegisteredNotAbleToSign - {epoch}")
75            }
76            Self::ReadyToSign { epoch } => {
77                write!(f, "ReadyToSign - {epoch}")
78            }
79        }
80    }
81}
82
83enum EpochStatus {
84    NewEpoch(Epoch),
85    Unchanged(TimePoint),
86}
87
88/// The state machine is responsible of the execution of the signer automate.
89pub struct StateMachine {
90    state: Mutex<SignerState>,
91    runner: Box<dyn Runner>,
92    interval: Duration,
93    metrics_service: Arc<MetricsService>,
94    logger: Logger,
95}
96
97impl StateMachine {
98    /// Create a new StateMachine instance.
99    pub fn new(
100        starting_state: SignerState,
101        runner: Box<dyn Runner>,
102        interval: Duration,
103        metrics_service: Arc<MetricsService>,
104        logger: Logger,
105    ) -> Self {
106        Self {
107            state: Mutex::new(starting_state),
108            runner,
109            interval,
110            metrics_service,
111            logger: logger.new_with_component_name::<Self>(),
112        }
113    }
114
115    /// Return the current state of the state machine.
116    pub async fn get_state(&self) -> SignerState {
117        self.state.lock().await.to_owned()
118    }
119
120    /// Launch the state machine until an error occurs or it is interrupted.
121    pub async fn run(&self) -> Result<(), RuntimeError> {
122        info!(self.logger, "Launching State Machine");
123        let mut interval = tokio::time::interval(self.interval);
124
125        loop {
126            interval.tick().await;
127            // Note: the "time" property in logs produced by our formatter (slog_bunyan) uses local
128            // time, so we must use it as well to avoid confusion.
129            let approximate_next_cycle_time = Local::now() + self.interval;
130
131            if let Err(e) = self.cycle().await {
132                e.write_to_log(&self.logger);
133                if e.is_critical() {
134                    return Err(e);
135                }
136            }
137
138            info!(
139                self.logger, "… Cycle finished";
140                "approximate_next_cycle_time" => %approximate_next_cycle_time.time().format("%H:%M:%S%.3f"),
141                "run_interval_in_ms" => self.interval.as_millis(),
142            );
143        }
144    }
145
146    /// Perform a cycle of the state machine.
147    pub async fn cycle(&self) -> Result<(), RuntimeError> {
148        let mut state = self.state.lock().await;
149        info!(
150            self.logger,
151            "================================================================================"
152        );
153        info!(self.logger, "New cycle: {}", *state);
154
155        self.metrics_service
156            .get_runtime_cycle_total_since_startup_counter()
157            .increment();
158
159        match state.deref() {
160            SignerState::Init => {
161                *state = self.transition_from_init_to_unregistered().await?;
162            }
163            SignerState::Unregistered { epoch } => {
164                if let EpochStatus::NewEpoch(new_epoch) = self.has_epoch_changed(*epoch).await? {
165                    info!(
166                        self.logger,
167                        "→ Epoch has changed, transiting to Unregistered"
168                    );
169                    *state = self
170                        .transition_from_unregistered_to_unregistered(new_epoch)
171                        .await?;
172                } else if let Some(epoch_settings) = self
173                    .runner
174                    .get_epoch_settings()
175                    .await
176                    .map_err(|e| RuntimeError::KeepState {
177                        message: format!("could not retrieve epoch settings at epoch {epoch:?}"),
178                        nested_error: Some(e),
179                    })?
180                {
181                    info!(self.logger, "→ Epoch settings found");
182                    if epoch_settings.epoch >= *epoch {
183                        info!(self.logger, "New Epoch found");
184                        info!(self.logger, " ⋅ Transiting to Registered");
185                        *state = self
186                            .transition_from_unregistered_to_one_of_registered_states(
187                                epoch_settings,
188                            )
189                            .await?;
190                    } else {
191                        info!(
192                            self.logger, " ⋅ Epoch settings found, but its epoch is behind the known epoch, waiting…";
193                            "epoch_settings" => ?epoch_settings,
194                            "known_epoch" => ?epoch,
195                        );
196                    }
197                } else {
198                    info!(self.logger, "→ No epoch settings found yet, waiting…");
199                }
200            }
201            SignerState::RegisteredNotAbleToSign { epoch } => {
202                if let EpochStatus::NewEpoch(new_epoch) = self.has_epoch_changed(*epoch).await? {
203                    info!(
204                        self.logger,
205                        " → New Epoch detected, transiting to Unregistered"
206                    );
207                    *state = self
208                        .transition_from_registered_not_able_to_sign_to_unregistered(new_epoch)
209                        .await?;
210                } else {
211                    info!(self.logger, " ⋅ Epoch has NOT changed, waiting…");
212                }
213            }
214
215            SignerState::ReadyToSign { epoch } => match self.has_epoch_changed(*epoch).await? {
216                EpochStatus::NewEpoch(new_epoch) => {
217                    info!(
218                        self.logger,
219                        "→ Epoch has changed, transiting to Unregistered"
220                    );
221                    *state = self
222                        .transition_from_ready_to_sign_to_unregistered(new_epoch)
223                        .await?;
224                }
225                EpochStatus::Unchanged(timepoint) => {
226                    let beacon_to_sign =
227                        self.runner
228                            .get_beacon_to_sign(timepoint)
229                            .await
230                            .map_err(|e| RuntimeError::KeepState {
231                                message: "could not fetch the beacon to sign".to_string(),
232                                nested_error: Some(e),
233                            })?;
234
235                    match beacon_to_sign {
236                        Some(beacon) => {
237                            info!(
238                                self.logger, "→ Epoch has NOT changed we can sign this beacon, transiting to ReadyToSign";
239                                "beacon_to_sign" => ?beacon,
240                            );
241                            *state = self
242                                .transition_from_ready_to_sign_to_ready_to_sign(*epoch, beacon)
243                                .await?;
244                        }
245                        None => {
246                            info!(self.logger, " ⋅ No beacon to sign, waiting…");
247                        }
248                    }
249                }
250            },
251        };
252
253        self.metrics_service
254            .get_runtime_cycle_success_since_startup_counter()
255            .increment();
256
257        Ok(())
258    }
259
260    /// Return the new epoch if the epoch is different than the given one, otherwise return the current time point.
261    async fn has_epoch_changed(&self, epoch: Epoch) -> Result<EpochStatus, RuntimeError> {
262        let current_time_point = self
263            .get_current_time_point("checking if epoch has changed")
264            .await?;
265
266        if current_time_point.epoch > epoch {
267            Ok(EpochStatus::NewEpoch(current_time_point.epoch))
268        } else {
269            Ok(EpochStatus::Unchanged(current_time_point))
270        }
271    }
272
273    async fn transition_from_unregistered_to_unregistered(
274        &self,
275        new_epoch: Epoch,
276    ) -> Result<SignerState, RuntimeError> {
277        self.update_era_checker(new_epoch, "unregistered → unregistered")
278            .await?;
279
280        Ok(SignerState::Unregistered { epoch: new_epoch })
281    }
282
283    async fn transition_from_init_to_unregistered(&self) -> Result<SignerState, RuntimeError> {
284        let current_epoch = self
285            .get_current_time_point("init → unregistered")
286            .await?
287            .epoch;
288        self.update_era_checker(current_epoch, "init → unregistered")
289            .await?;
290
291        Ok(SignerState::Unregistered {
292            epoch: current_epoch,
293        })
294    }
295
296    /// Launch the transition process from the `Unregistered` to `ReadyToSign` or `RegisteredNotAbleToSign` state.
297    async fn transition_from_unregistered_to_one_of_registered_states(
298        &self,
299        epoch_settings: SignerEpochSettings,
300    ) -> Result<SignerState, RuntimeError> {
301        self.metrics_service
302            .get_signer_registration_total_since_startup_counter()
303            .increment();
304
305        let time_point = self
306            .get_current_time_point("unregistered → registered")
307            .await?;
308        let epoch = time_point.epoch;
309        self.runner.update_stake_distribution(epoch)
310            .await
311            .map_err(|e| RuntimeError::KeepState {
312                message: format!("Could not update stake distribution in 'unregistered → registered' phase for epoch {epoch:?}."),
313                nested_error: Some(e),
314            })?;
315
316        self.runner
317            .inform_epoch_settings(epoch_settings)
318            .await
319            .map_err(|e| RuntimeError::KeepState {
320                message: format!(
321                    "Could not register epoch information in 'unregistered → registered' phase for epoch {epoch:?}."
322                ),
323                nested_error: Some(e),
324            })?;
325
326        fn handle_registration_result(
327            register_result: Result<(), Error>,
328            epoch: Epoch,
329        ) -> Result<Option<SignerState>, RuntimeError> {
330            if let Err(e) = register_result {
331                if let Some(AggregatorClientError::RegistrationRoundNotYetOpened(_)) =
332                    e.downcast_ref::<AggregatorClientError>()
333                {
334                    Ok(Some(SignerState::Unregistered { epoch }))
335                } else if e.downcast_ref::<ProtocolInitializerError>().is_some() {
336                    Err(RuntimeError::Critical { message: format!("Could not register to aggregator in 'unregistered → registered' phase for epoch {epoch:?}."), nested_error: Some(e) })
337                } else {
338                    Err(RuntimeError::KeepState { message: format!("Could not register to aggregator in 'unregistered → registered' phase for epoch {epoch:?}."), nested_error: Some(e) })
339                }
340            } else {
341                Ok(None)
342            }
343        }
344
345        let register_result = self.runner.register_signer_to_aggregator().await;
346        let next_state_found = handle_registration_result(register_result, epoch)?;
347
348        self.metrics_service
349            .get_signer_registration_success_since_startup_counter()
350            .increment();
351
352        if let Some(state) = next_state_found {
353            return Ok(state);
354        }
355
356        self.metrics_service
357            .get_signer_registration_success_last_epoch_gauge()
358            .record(epoch);
359
360        self.runner
361            .upkeep(epoch)
362            .await
363            .map_err(|e| RuntimeError::KeepState {
364                message: "Failed to upkeep signer in 'unregistered → registered' phase".to_string(),
365                nested_error: Some(e),
366            })?;
367
368        match self
369            .runner
370            .can_sign_current_epoch()
371            .await
372            .map_err(|e| RuntimeError::KeepState {
373                message: "Failed to check if signer can sign in the current epoch in 'unregistered → ?' phase".to_string(),
374                nested_error: Some(e),
375            })? {
376            true => Ok(SignerState::ReadyToSign { epoch }),
377            false => Ok(SignerState::RegisteredNotAbleToSign { epoch }),
378        }
379    }
380
381    async fn transition_from_registered_not_able_to_sign_to_unregistered(
382        &self,
383        epoch: Epoch,
384    ) -> Result<SignerState, RuntimeError> {
385        self.update_era_checker(epoch, "registered not able to sign → unregistered")
386            .await?;
387
388        Ok(SignerState::Unregistered { epoch })
389    }
390
391    async fn transition_from_ready_to_sign_to_unregistered(
392        &self,
393        epoch: Epoch,
394    ) -> Result<SignerState, RuntimeError> {
395        self.update_era_checker(epoch, "ready to sign → unregistered")
396            .await?;
397
398        Ok(SignerState::Unregistered { epoch })
399    }
400
401    /// Launch the transition process from the `ReadyToSign` to the `ReadyToSign` state.
402    async fn transition_from_ready_to_sign_to_ready_to_sign(
403        &self,
404        current_epoch: Epoch,
405        beacon_to_sign: BeaconToSign,
406    ) -> Result<SignerState, RuntimeError> {
407        let (retrieval_epoch, next_retrieval_epoch) = (
408            current_epoch.offset_to_signer_retrieval_epoch()?,
409            current_epoch.offset_to_next_signer_retrieval_epoch(),
410        );
411
412        debug!(
413            self.logger, ">> transition_from_ready_to_sign_to_ready_to_sign";
414            "current_epoch" => ?current_epoch,
415            "retrieval_epoch" => ?retrieval_epoch,
416            "next_retrieval_epoch" => ?next_retrieval_epoch,
417        );
418
419        self.metrics_service
420            .get_signature_registration_total_since_startup_counter()
421            .increment();
422
423        let message = self
424            .runner
425            .compute_message(&beacon_to_sign.signed_entity_type)
426            .await
427            .map_err(|e| RuntimeError::KeepState {
428                message: format!("Could not compute message during 'ready to sign → ready to sign' phase (current epoch {current_epoch:?})"),
429                nested_error: Some(e)
430            })?;
431
432        self.runner.compute_publish_single_signature(&beacon_to_sign, &message)
433            .await
434            .map_err(|e| RuntimeError::KeepState {
435                message: format!("Could not compute and publish single signature during 'ready to sign → ready to sign' phase (current epoch {current_epoch:?})"),
436                nested_error: Some(e)
437            })?;
438
439        self.metrics_service
440            .get_signature_registration_success_since_startup_counter()
441            .increment();
442        self.metrics_service
443            .get_signature_registration_success_last_epoch_gauge()
444            .record(current_epoch);
445
446        Ok(SignerState::ReadyToSign {
447            epoch: current_epoch,
448        })
449    }
450
451    async fn get_current_time_point(&self, context: &str) -> Result<TimePoint, RuntimeError> {
452        let current_time_point =
453            self.runner
454                .get_current_time_point()
455                .await
456                .map_err(|e| RuntimeError::KeepState {
457                    message: format!(
458                        "Could not retrieve current time point in context '{context}'."
459                    ),
460                    nested_error: Some(e),
461                })?;
462
463        Ok(current_time_point)
464    }
465
466    async fn update_era_checker(&self, epoch: Epoch, context: &str) -> Result<(), RuntimeError> {
467        self.runner
468            .update_era_checker(epoch)
469            .await
470            .map_err(|e| RuntimeError::Critical {
471                message: format!(
472                    "Could not update Era checker with context '{context}' for epoch {epoch:?}"
473                ),
474                nested_error: Some(e),
475            })
476    }
477}
478
479#[cfg(test)]
480mod tests {
481    use anyhow::anyhow;
482    use chrono::DateTime;
483    use mockall::predicate;
484
485    use mithril_common::entities::{ChainPoint, Epoch, ProtocolMessage, SignedEntityType};
486    use mithril_common::test_utils::fake_data;
487
488    use crate::runtime::runner::MockSignerRunner;
489    use crate::services::AggregatorClientError;
490    use crate::test_tools::TestLogger;
491
492    use super::*;
493
494    fn init_state_machine(init_state: SignerState, runner: MockSignerRunner) -> StateMachine {
495        let logger = TestLogger::stdout();
496        let metrics_service = Arc::new(MetricsService::new(logger.clone()).unwrap());
497        StateMachine {
498            state: init_state.into(),
499            runner: Box::new(runner),
500            interval: Duration::from_millis(100),
501            metrics_service,
502            logger,
503        }
504    }
505
506    #[tokio::test]
507    async fn unregistered_epoch_settings_not_found() {
508        let mut runner = MockSignerRunner::new();
509        runner
510            .expect_get_epoch_settings()
511            .once()
512            .returning(|| Ok(None));
513        runner
514            .expect_get_current_time_point()
515            .once()
516            .returning(|| Ok(TimePoint::dummy()));
517        let state_machine = init_state_machine(
518            SignerState::Unregistered {
519                epoch: TimePoint::dummy().epoch,
520            },
521            runner,
522        );
523        state_machine
524            .cycle()
525            .await
526            .expect("Cycling the state machine should not fail");
527
528        assert_eq!(
529            SignerState::Unregistered {
530                epoch: TimePoint::dummy().epoch
531            },
532            state_machine.get_state().await
533        );
534    }
535
536    #[tokio::test]
537    async fn unregistered_epoch_settings_behind_known_epoch() {
538        let mut runner = MockSignerRunner::new();
539        let epoch_settings = SignerEpochSettings {
540            epoch: Epoch(3),
541            registration_protocol_parameters: fake_data::protocol_parameters(),
542            current_signers: vec![],
543            next_signers: vec![],
544            cardano_transactions_signing_config: None,
545            next_cardano_transactions_signing_config: None,
546        };
547        let known_epoch = Epoch(4);
548        runner
549            .expect_get_epoch_settings()
550            .once()
551            .returning(move || Ok(Some(epoch_settings.to_owned())));
552        runner.expect_get_current_time_point().once().returning(|| {
553            Ok(TimePoint {
554                epoch: Epoch(4),
555                ..TimePoint::dummy()
556            })
557        });
558        let state_machine =
559            init_state_machine(SignerState::Unregistered { epoch: known_epoch }, runner);
560        state_machine
561            .cycle()
562            .await
563            .expect("Cycling the state machine should not fail");
564
565        assert_eq!(
566            SignerState::Unregistered { epoch: known_epoch },
567            state_machine.get_state().await
568        );
569    }
570
571    #[tokio::test]
572    async fn unregistered_to_registered_not_able_to_sign() {
573        let mut runner = MockSignerRunner::new();
574        runner.expect_upkeep().returning(|_| Ok(())).once();
575        runner
576            .expect_get_epoch_settings()
577            .once()
578            .returning(|| Ok(Some(SignerEpochSettings::dummy())));
579
580        runner
581            .expect_inform_epoch_settings()
582            .with(predicate::eq(SignerEpochSettings::dummy()))
583            .once()
584            .returning(|_| Ok(()));
585
586        runner
587            .expect_get_current_time_point()
588            .times(2)
589            .returning(|| Ok(TimePoint::dummy()));
590        runner
591            .expect_update_stake_distribution()
592            .once()
593            .returning(|_| Ok(()));
594        runner
595            .expect_register_signer_to_aggregator()
596            .once()
597            .returning(|| Ok(()));
598
599        runner
600            .expect_can_sign_current_epoch()
601            .once()
602            .returning(|| Ok(false));
603
604        let state_machine = init_state_machine(
605            SignerState::Unregistered {
606                epoch: TimePoint::dummy().epoch,
607            },
608            runner,
609        );
610
611        state_machine
612            .cycle()
613            .await
614            .expect("Cycling the state machine should not fail");
615
616        if let SignerState::RegisteredNotAbleToSign { epoch: _ } = state_machine.get_state().await {
617        } else {
618            panic!(
619                "state machine did not return a RegisteredNotAbleToSign state but {:?}",
620                state_machine.get_state().await
621            );
622        }
623    }
624
625    #[tokio::test]
626    async fn unregistered_to_ready_to_sign() {
627        let mut runner = MockSignerRunner::new();
628        runner.expect_upkeep().returning(|_| Ok(())).once();
629        runner
630            .expect_get_epoch_settings()
631            .once()
632            .returning(|| Ok(Some(SignerEpochSettings::dummy())));
633
634        runner
635            .expect_inform_epoch_settings()
636            .with(predicate::eq(SignerEpochSettings::dummy()))
637            .once()
638            .returning(|_| Ok(()));
639
640        runner
641            .expect_get_current_time_point()
642            .times(2)
643            .returning(|| Ok(TimePoint::dummy()));
644        runner
645            .expect_update_stake_distribution()
646            .once()
647            .returning(|_| Ok(()));
648        runner
649            .expect_register_signer_to_aggregator()
650            .once()
651            .returning(|| Ok(()));
652
653        runner
654            .expect_can_sign_current_epoch()
655            .once()
656            .returning(|| Ok(true));
657
658        let state_machine = init_state_machine(
659            SignerState::Unregistered {
660                epoch: TimePoint::dummy().epoch,
661            },
662            runner,
663        );
664
665        state_machine
666            .cycle()
667            .await
668            .expect("Cycling the state machine should not fail");
669
670        assert_eq!(
671            SignerState::ReadyToSign {
672                epoch: TimePoint::dummy().epoch,
673            },
674            state_machine.get_state().await
675        );
676
677        let metrics_service = state_machine.metrics_service;
678        let success_since_startup =
679            metrics_service.get_runtime_cycle_success_since_startup_counter();
680        assert_eq!(1, success_since_startup.get());
681    }
682
683    #[tokio::test]
684    async fn unregistered_to_ready_to_sign_counter() {
685        let mut runner = MockSignerRunner::new();
686
687        runner
688            .expect_get_epoch_settings()
689            .once()
690            .returning(|| Ok(Some(SignerEpochSettings::dummy())));
691
692        runner
693            .expect_inform_epoch_settings()
694            .with(predicate::eq(SignerEpochSettings::dummy()))
695            .once()
696            .returning(|_| Ok(()));
697
698        runner
699            .expect_get_current_time_point()
700            .times(2)
701            .returning(|| Ok(TimePoint::dummy()));
702        runner
703            .expect_update_stake_distribution()
704            .once()
705            .returning(|_| Ok(()));
706        runner
707            .expect_register_signer_to_aggregator()
708            .once()
709            .returning(|| {
710                Err(AggregatorClientError::RegistrationRoundNotYetOpened(
711                    anyhow!("Not yet opened"),
712                ))?
713            });
714
715        runner.expect_upkeep().never();
716        runner.expect_can_sign_current_epoch().never();
717
718        let state_machine = init_state_machine(
719            SignerState::Unregistered {
720                epoch: TimePoint::dummy().epoch,
721            },
722            runner,
723        );
724
725        state_machine
726            .cycle()
727            .await
728            .expect("Cycling the state machine should not fail");
729
730        assert_eq!(
731            SignerState::Unregistered {
732                epoch: TimePoint::dummy().epoch,
733            },
734            state_machine.get_state().await
735        );
736
737        let metrics_service = state_machine.metrics_service;
738        assert_eq!(
739            1,
740            metrics_service
741                .get_signer_registration_success_since_startup_counter()
742                .get()
743        );
744
745        assert_eq!(
746            0 as f64,
747            metrics_service
748                .get_signer_registration_success_last_epoch_gauge()
749                .get()
750        );
751
752        assert_eq!(
753            1,
754            metrics_service
755                .get_runtime_cycle_total_since_startup_counter()
756                .get()
757        );
758    }
759
760    #[tokio::test]
761    async fn registered_not_able_to_sign_to_unregistered() {
762        let mut runner = MockSignerRunner::new();
763        runner
764            .expect_get_current_time_point()
765            .once()
766            .returning(|| Ok(TimePoint::new(10, 100, ChainPoint::dummy())));
767        runner
768            .expect_update_era_checker()
769            .once()
770            .returning(|_e: Epoch| Ok(()));
771
772        let state_machine = init_state_machine(
773            SignerState::RegisteredNotAbleToSign { epoch: Epoch(0) },
774            runner,
775        );
776
777        state_machine
778            .cycle()
779            .await
780            .expect("Cycling the state machine should not fail");
781        assert_eq!(
782            SignerState::Unregistered { epoch: Epoch(10) },
783            state_machine.get_state().await
784        );
785    }
786
787    #[tokio::test]
788    async fn registered_not_able_to_sign_to_registered_not_able_to_sign() {
789        let mut runner = MockSignerRunner::new();
790        runner
791            .expect_get_current_time_point()
792            .once()
793            .returning(|| Ok(TimePoint::new(10, 100, ChainPoint::dummy())));
794
795        let state_machine = init_state_machine(
796            SignerState::RegisteredNotAbleToSign { epoch: Epoch(10) },
797            runner,
798        );
799
800        state_machine
801            .cycle()
802            .await
803            .expect("Cycling the state machine should not fail");
804        assert_eq!(
805            SignerState::RegisteredNotAbleToSign { epoch: Epoch(10) },
806            state_machine.get_state().await
807        );
808    }
809
810    #[tokio::test]
811    async fn ready_to_sign_to_unregistered() {
812        let mut runner = MockSignerRunner::new();
813        runner
814            .expect_get_current_time_point()
815            .once()
816            .returning(|| Ok(TimePoint::new(10, 100, ChainPoint::dummy())));
817        runner
818            .expect_update_era_checker()
819            .once()
820            .returning(|_e: Epoch| Ok(()));
821
822        let state_machine =
823            init_state_machine(SignerState::ReadyToSign { epoch: Epoch(0) }, runner);
824
825        state_machine
826            .cycle()
827            .await
828            .expect("Cycling the state machine should not fail");
829        assert_eq!(
830            SignerState::Unregistered { epoch: Epoch(10) },
831            state_machine.get_state().await
832        );
833    }
834
835    #[tokio::test]
836    async fn ready_to_sign_to_ready_to_sign_when_there_is_a_beacon_to_sign() {
837        let time_point = TimePoint::dummy();
838        let beacon_to_sign = BeaconToSign {
839            epoch: time_point.epoch,
840            signed_entity_type: SignedEntityType::MithrilStakeDistribution(time_point.epoch),
841            initiated_at: DateTime::default(),
842        };
843        let beacon_to_sign_clone = beacon_to_sign.clone();
844        let current_epoch = time_point.epoch;
845
846        let mut runner = MockSignerRunner::new();
847        runner
848            .expect_get_current_time_point()
849            .once()
850            .returning(move || Ok(time_point.to_owned()));
851        runner
852            .expect_get_beacon_to_sign()
853            .once()
854            .returning(move |_| Ok(Some(beacon_to_sign_clone.clone())));
855        runner
856            .expect_compute_message()
857            .once()
858            .returning(|_| Ok(ProtocolMessage::new()));
859        runner
860            .expect_compute_publish_single_signature()
861            .once()
862            .returning(|_, _| Ok(()));
863
864        let state_machine = init_state_machine(
865            SignerState::ReadyToSign {
866                epoch: current_epoch,
867            },
868            runner,
869        );
870        state_machine
871            .cycle()
872            .await
873            .expect("Cycling the state machine should not fail");
874
875        assert_eq!(
876            SignerState::ReadyToSign {
877                epoch: current_epoch
878            },
879            state_machine.get_state().await,
880            "state machine did not return a ReadyToSign state but {:?}",
881            state_machine.get_state().await
882        );
883    }
884
885    #[tokio::test]
886    async fn ready_to_sign_to_ready_to_sign_when_there_no_beacon_to_sign() {
887        let time_point = TimePoint::dummy();
888        let current_epoch = time_point.epoch;
889
890        let mut runner = MockSignerRunner::new();
891        runner
892            .expect_get_current_time_point()
893            .once()
894            .returning(move || Ok(time_point.to_owned()));
895        runner
896            .expect_get_beacon_to_sign()
897            .once()
898            .returning(move |_| Ok(None));
899
900        let state_machine = init_state_machine(
901            SignerState::ReadyToSign {
902                epoch: current_epoch,
903            },
904            runner,
905        );
906        state_machine
907            .cycle()
908            .await
909            .expect("Cycling the state machine should not fail");
910
911        assert_eq!(
912            SignerState::ReadyToSign {
913                epoch: current_epoch
914            },
915            state_machine.get_state().await,
916            "state machine did not return a ReadyToSign state but {:?}",
917            state_machine.get_state().await
918        );
919    }
920}