mithril_signer/runtime/
state_machine.rs

1use anyhow::Error;
2use chrono::Local;
3use slog::{Logger, debug, info};
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    MetricsService,
15    entities::{BeaconToSign, SignerEpochSettings},
16    services::AggregatorClientError,
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.transition_from_unregistered_to_unregistered(new_epoch).await?;
170                } else if let Some(epoch_settings) = self
171                    .runner
172                    .get_epoch_settings()
173                    .await
174                    .map_err(|e| RuntimeError::KeepState {
175                        message: format!("could not retrieve epoch settings at epoch {epoch:?}"),
176                        nested_error: Some(e),
177                    })?
178                {
179                    info!(self.logger, "→ Epoch settings found");
180                    if epoch_settings.epoch >= *epoch {
181                        info!(self.logger, "New Epoch found");
182                        info!(self.logger, " ⋅ Transiting to Registered");
183                        *state = self
184                            .transition_from_unregistered_to_one_of_registered_states(
185                                epoch_settings,
186                            )
187                            .await?;
188                    } else {
189                        info!(
190                            self.logger, " ⋅ Epoch settings found, but its epoch is behind the known epoch, waiting…";
191                            "epoch_settings" => ?epoch_settings,
192                            "known_epoch" => ?epoch,
193                        );
194                    }
195                } else {
196                    info!(self.logger, "→ No epoch settings found yet, waiting…");
197                }
198            }
199            SignerState::RegisteredNotAbleToSign { epoch } => {
200                if let EpochStatus::NewEpoch(new_epoch) = self.has_epoch_changed(*epoch).await? {
201                    info!(
202                        self.logger,
203                        " → New Epoch detected, transiting to Unregistered"
204                    );
205                    *state = self
206                        .transition_from_registered_not_able_to_sign_to_unregistered(new_epoch)
207                        .await?;
208                } else {
209                    info!(self.logger, " ⋅ Epoch has NOT changed, waiting…");
210                }
211            }
212
213            SignerState::ReadyToSign { epoch } => match self.has_epoch_changed(*epoch).await? {
214                EpochStatus::NewEpoch(new_epoch) => {
215                    info!(
216                        self.logger,
217                        "→ Epoch has changed, transiting to Unregistered"
218                    );
219                    *state = self.transition_from_ready_to_sign_to_unregistered(new_epoch).await?;
220                }
221                EpochStatus::Unchanged(timepoint) => {
222                    let beacon_to_sign =
223                        self.runner.get_beacon_to_sign(timepoint).await.map_err(|e| {
224                            RuntimeError::KeepState {
225                                message: "could not fetch the beacon to sign".to_string(),
226                                nested_error: Some(e),
227                            }
228                        })?;
229
230                    match beacon_to_sign {
231                        Some(beacon) => {
232                            info!(
233                                self.logger, "→ Epoch has NOT changed we can sign this beacon, transiting to ReadyToSign";
234                                "beacon_to_sign" => ?beacon,
235                            );
236                            *state = self
237                                .transition_from_ready_to_sign_to_ready_to_sign(*epoch, beacon)
238                                .await?;
239                        }
240                        None => {
241                            info!(self.logger, " ⋅ No beacon to sign, waiting…");
242                        }
243                    }
244                }
245            },
246        };
247
248        self.metrics_service
249            .get_runtime_cycle_success_since_startup_counter()
250            .increment();
251
252        Ok(())
253    }
254
255    /// Return the new epoch if the epoch is different than the given one, otherwise return the current time point.
256    async fn has_epoch_changed(&self, epoch: Epoch) -> Result<EpochStatus, RuntimeError> {
257        let current_time_point =
258            self.get_current_time_point("checking if epoch has changed").await?;
259
260        if current_time_point.epoch > epoch {
261            Ok(EpochStatus::NewEpoch(current_time_point.epoch))
262        } else {
263            Ok(EpochStatus::Unchanged(current_time_point))
264        }
265    }
266
267    async fn transition_from_unregistered_to_unregistered(
268        &self,
269        new_epoch: Epoch,
270    ) -> Result<SignerState, RuntimeError> {
271        self.update_era_checker(new_epoch, "unregistered → unregistered")
272            .await?;
273
274        Ok(SignerState::Unregistered { epoch: new_epoch })
275    }
276
277    async fn transition_from_init_to_unregistered(&self) -> Result<SignerState, RuntimeError> {
278        let current_epoch = self.get_current_time_point("init → unregistered").await?.epoch;
279        self.update_era_checker(current_epoch, "init → unregistered").await?;
280
281        Ok(SignerState::Unregistered {
282            epoch: current_epoch,
283        })
284    }
285
286    /// Launch the transition process from the `Unregistered` to `ReadyToSign` or `RegisteredNotAbleToSign` state.
287    async fn transition_from_unregistered_to_one_of_registered_states(
288        &self,
289        epoch_settings: SignerEpochSettings,
290    ) -> Result<SignerState, RuntimeError> {
291        self.metrics_service
292            .get_signer_registration_total_since_startup_counter()
293            .increment();
294
295        let time_point = self.get_current_time_point("unregistered → registered").await?;
296        let epoch = time_point.epoch;
297        self.runner.update_stake_distribution(epoch)
298            .await
299            .map_err(|e| RuntimeError::KeepState {
300                message: format!("Could not update stake distribution in 'unregistered → registered' phase for epoch {epoch:?}."),
301                nested_error: Some(e),
302            })?;
303
304        self.runner
305            .inform_epoch_settings(epoch_settings)
306            .await
307            .map_err(|e| RuntimeError::KeepState {
308                message: format!(
309                    "Could not register epoch information in 'unregistered → registered' phase for epoch {epoch:?}."
310                ),
311                nested_error: Some(e),
312            })?;
313
314        fn handle_registration_result(
315            register_result: Result<(), Error>,
316            epoch: Epoch,
317        ) -> Result<Option<SignerState>, RuntimeError> {
318            if let Err(e) = register_result {
319                if let Some(AggregatorClientError::RegistrationRoundNotYetOpened(_)) =
320                    e.downcast_ref::<AggregatorClientError>()
321                {
322                    Ok(Some(SignerState::Unregistered { epoch }))
323                } else if e.downcast_ref::<ProtocolInitializerError>().is_some() {
324                    Err(RuntimeError::Critical {
325                        message: format!(
326                            "Could not register to aggregator in 'unregistered → registered' phase for epoch {epoch:?}."
327                        ),
328                        nested_error: Some(e),
329                    })
330                } else {
331                    Err(RuntimeError::KeepState {
332                        message: format!(
333                            "Could not register to aggregator in 'unregistered → registered' phase for epoch {epoch:?}."
334                        ),
335                        nested_error: Some(e),
336                    })
337                }
338            } else {
339                Ok(None)
340            }
341        }
342
343        let register_result = self.runner.register_signer_to_aggregator().await;
344        let next_state_found = handle_registration_result(register_result, epoch)?;
345
346        self.metrics_service
347            .get_signer_registration_success_since_startup_counter()
348            .increment();
349
350        if let Some(state) = next_state_found {
351            return Ok(state);
352        }
353
354        self.metrics_service
355            .get_signer_registration_success_last_epoch_gauge()
356            .record(epoch);
357
358        self.runner.upkeep(epoch).await.map_err(|e| RuntimeError::KeepState {
359            message: "Failed to upkeep signer in 'unregistered → registered' phase".to_string(),
360            nested_error: Some(e),
361        })?;
362
363        match self
364            .runner
365            .can_sign_current_epoch()
366            .await
367            .map_err(|e| RuntimeError::KeepState {
368                message: "Failed to check if signer can sign in the current epoch in 'unregistered → ?' phase".to_string(),
369                nested_error: Some(e),
370            })? {
371            true => Ok(SignerState::ReadyToSign { epoch }),
372            false => Ok(SignerState::RegisteredNotAbleToSign { epoch }),
373        }
374    }
375
376    async fn transition_from_registered_not_able_to_sign_to_unregistered(
377        &self,
378        epoch: Epoch,
379    ) -> Result<SignerState, RuntimeError> {
380        self.update_era_checker(epoch, "registered not able to sign → unregistered")
381            .await?;
382
383        Ok(SignerState::Unregistered { epoch })
384    }
385
386    async fn transition_from_ready_to_sign_to_unregistered(
387        &self,
388        epoch: Epoch,
389    ) -> Result<SignerState, RuntimeError> {
390        self.update_era_checker(epoch, "ready to sign → unregistered").await?;
391
392        Ok(SignerState::Unregistered { epoch })
393    }
394
395    /// Launch the transition process from the `ReadyToSign` to the `ReadyToSign` state.
396    async fn transition_from_ready_to_sign_to_ready_to_sign(
397        &self,
398        current_epoch: Epoch,
399        beacon_to_sign: BeaconToSign,
400    ) -> Result<SignerState, RuntimeError> {
401        let (retrieval_epoch, next_retrieval_epoch) = (
402            current_epoch.offset_to_signer_retrieval_epoch()?,
403            current_epoch.offset_to_next_signer_retrieval_epoch(),
404        );
405
406        debug!(
407            self.logger, ">> transition_from_ready_to_sign_to_ready_to_sign";
408            "current_epoch" => ?current_epoch,
409            "retrieval_epoch" => ?retrieval_epoch,
410            "next_retrieval_epoch" => ?next_retrieval_epoch,
411        );
412
413        self.metrics_service
414            .get_signature_registration_total_since_startup_counter()
415            .increment();
416
417        let message = self
418            .runner
419            .compute_message(&beacon_to_sign.signed_entity_type)
420            .await
421            .map_err(|e| RuntimeError::KeepState {
422                message: format!("Could not compute message during 'ready to sign → ready to sign' phase (current epoch {current_epoch:?})"),
423                nested_error: Some(e)
424            })?;
425
426        self.runner.compute_publish_single_signature(&beacon_to_sign, &message)
427            .await
428            .map_err(|e| RuntimeError::KeepState {
429                message: format!("Could not compute and publish single signature during 'ready to sign → ready to sign' phase (current epoch {current_epoch:?})"),
430                nested_error: Some(e)
431            })?;
432
433        self.metrics_service
434            .get_signature_registration_success_since_startup_counter()
435            .increment();
436        self.metrics_service
437            .get_signature_registration_success_last_epoch_gauge()
438            .record(current_epoch);
439
440        Ok(SignerState::ReadyToSign {
441            epoch: current_epoch,
442        })
443    }
444
445    async fn get_current_time_point(&self, context: &str) -> Result<TimePoint, RuntimeError> {
446        let current_time_point =
447            self.runner
448                .get_current_time_point()
449                .await
450                .map_err(|e| RuntimeError::KeepState {
451                    message: format!(
452                        "Could not retrieve current time point in context '{context}'."
453                    ),
454                    nested_error: Some(e),
455                })?;
456
457        Ok(current_time_point)
458    }
459
460    async fn update_era_checker(&self, epoch: Epoch, context: &str) -> Result<(), RuntimeError> {
461        self.runner
462            .update_era_checker(epoch)
463            .await
464            .map_err(|e| RuntimeError::Critical {
465                message: format!(
466                    "Could not update Era checker with context '{context}' for epoch {epoch:?}"
467                ),
468                nested_error: Some(e),
469            })
470    }
471}
472
473#[cfg(test)]
474mod tests {
475    use anyhow::anyhow;
476    use chrono::DateTime;
477    use mockall::predicate;
478
479    use mithril_common::entities::{ChainPoint, Epoch, ProtocolMessage, SignedEntityType};
480    use mithril_common::test_utils::fake_data;
481
482    use crate::runtime::runner::MockSignerRunner;
483    use crate::services::AggregatorClientError;
484    use crate::test_tools::TestLogger;
485
486    use super::*;
487
488    fn init_state_machine(init_state: SignerState, runner: MockSignerRunner) -> StateMachine {
489        let logger = TestLogger::stdout();
490        let metrics_service = Arc::new(MetricsService::new(logger.clone()).unwrap());
491        StateMachine {
492            state: init_state.into(),
493            runner: Box::new(runner),
494            interval: Duration::from_millis(100),
495            metrics_service,
496            logger,
497        }
498    }
499
500    #[tokio::test]
501    async fn unregistered_epoch_settings_not_found() {
502        let mut runner = MockSignerRunner::new();
503        runner.expect_get_epoch_settings().once().returning(|| Ok(None));
504        runner
505            .expect_get_current_time_point()
506            .once()
507            .returning(|| Ok(TimePoint::dummy()));
508        let state_machine = init_state_machine(
509            SignerState::Unregistered {
510                epoch: TimePoint::dummy().epoch,
511            },
512            runner,
513        );
514        state_machine
515            .cycle()
516            .await
517            .expect("Cycling the state machine should not fail");
518
519        assert_eq!(
520            SignerState::Unregistered {
521                epoch: TimePoint::dummy().epoch
522            },
523            state_machine.get_state().await
524        );
525    }
526
527    #[tokio::test]
528    async fn unregistered_epoch_settings_behind_known_epoch() {
529        let mut runner = MockSignerRunner::new();
530        let epoch_settings = SignerEpochSettings {
531            epoch: Epoch(3),
532            registration_protocol_parameters: fake_data::protocol_parameters(),
533            current_signers: vec![],
534            next_signers: vec![],
535            cardano_transactions_signing_config: None,
536            next_cardano_transactions_signing_config: None,
537        };
538        let known_epoch = Epoch(4);
539        runner
540            .expect_get_epoch_settings()
541            .once()
542            .returning(move || Ok(Some(epoch_settings.to_owned())));
543        runner.expect_get_current_time_point().once().returning(|| {
544            Ok(TimePoint {
545                epoch: Epoch(4),
546                ..TimePoint::dummy()
547            })
548        });
549        let state_machine =
550            init_state_machine(SignerState::Unregistered { epoch: known_epoch }, runner);
551        state_machine
552            .cycle()
553            .await
554            .expect("Cycling the state machine should not fail");
555
556        assert_eq!(
557            SignerState::Unregistered { epoch: known_epoch },
558            state_machine.get_state().await
559        );
560    }
561
562    #[tokio::test]
563    async fn unregistered_to_registered_not_able_to_sign() {
564        let mut runner = MockSignerRunner::new();
565        runner.expect_upkeep().returning(|_| Ok(())).once();
566        runner
567            .expect_get_epoch_settings()
568            .once()
569            .returning(|| Ok(Some(SignerEpochSettings::dummy())));
570
571        runner
572            .expect_inform_epoch_settings()
573            .with(predicate::eq(SignerEpochSettings::dummy()))
574            .once()
575            .returning(|_| Ok(()));
576
577        runner
578            .expect_get_current_time_point()
579            .times(2)
580            .returning(|| Ok(TimePoint::dummy()));
581        runner.expect_update_stake_distribution().once().returning(|_| Ok(()));
582        runner
583            .expect_register_signer_to_aggregator()
584            .once()
585            .returning(|| Ok(()));
586
587        runner.expect_can_sign_current_epoch().once().returning(|| Ok(false));
588
589        let state_machine = init_state_machine(
590            SignerState::Unregistered {
591                epoch: TimePoint::dummy().epoch,
592            },
593            runner,
594        );
595
596        state_machine
597            .cycle()
598            .await
599            .expect("Cycling the state machine should not fail");
600
601        if let SignerState::RegisteredNotAbleToSign { epoch: _ } = state_machine.get_state().await {
602        } else {
603            panic!(
604                "state machine did not return a RegisteredNotAbleToSign state but {:?}",
605                state_machine.get_state().await
606            );
607        }
608    }
609
610    #[tokio::test]
611    async fn unregistered_to_ready_to_sign() {
612        let mut runner = MockSignerRunner::new();
613        runner.expect_upkeep().returning(|_| Ok(())).once();
614        runner
615            .expect_get_epoch_settings()
616            .once()
617            .returning(|| Ok(Some(SignerEpochSettings::dummy())));
618
619        runner
620            .expect_inform_epoch_settings()
621            .with(predicate::eq(SignerEpochSettings::dummy()))
622            .once()
623            .returning(|_| Ok(()));
624
625        runner
626            .expect_get_current_time_point()
627            .times(2)
628            .returning(|| Ok(TimePoint::dummy()));
629        runner.expect_update_stake_distribution().once().returning(|_| Ok(()));
630        runner
631            .expect_register_signer_to_aggregator()
632            .once()
633            .returning(|| Ok(()));
634
635        runner.expect_can_sign_current_epoch().once().returning(|| Ok(true));
636
637        let state_machine = init_state_machine(
638            SignerState::Unregistered {
639                epoch: TimePoint::dummy().epoch,
640            },
641            runner,
642        );
643
644        state_machine
645            .cycle()
646            .await
647            .expect("Cycling the state machine should not fail");
648
649        assert_eq!(
650            SignerState::ReadyToSign {
651                epoch: TimePoint::dummy().epoch,
652            },
653            state_machine.get_state().await
654        );
655
656        let metrics_service = state_machine.metrics_service;
657        let success_since_startup =
658            metrics_service.get_runtime_cycle_success_since_startup_counter();
659        assert_eq!(1, success_since_startup.get());
660    }
661
662    #[tokio::test]
663    async fn unregistered_to_ready_to_sign_counter() {
664        let mut runner = MockSignerRunner::new();
665
666        runner
667            .expect_get_epoch_settings()
668            .once()
669            .returning(|| Ok(Some(SignerEpochSettings::dummy())));
670
671        runner
672            .expect_inform_epoch_settings()
673            .with(predicate::eq(SignerEpochSettings::dummy()))
674            .once()
675            .returning(|_| Ok(()));
676
677        runner
678            .expect_get_current_time_point()
679            .times(2)
680            .returning(|| Ok(TimePoint::dummy()));
681        runner.expect_update_stake_distribution().once().returning(|_| Ok(()));
682        runner.expect_register_signer_to_aggregator().once().returning(|| {
683            Err(AggregatorClientError::RegistrationRoundNotYetOpened(
684                anyhow!("Not yet opened"),
685            ))?
686        });
687
688        runner.expect_upkeep().never();
689        runner.expect_can_sign_current_epoch().never();
690
691        let state_machine = init_state_machine(
692            SignerState::Unregistered {
693                epoch: TimePoint::dummy().epoch,
694            },
695            runner,
696        );
697
698        state_machine
699            .cycle()
700            .await
701            .expect("Cycling the state machine should not fail");
702
703        assert_eq!(
704            SignerState::Unregistered {
705                epoch: TimePoint::dummy().epoch,
706            },
707            state_machine.get_state().await
708        );
709
710        let metrics_service = state_machine.metrics_service;
711        assert_eq!(
712            1,
713            metrics_service
714                .get_signer_registration_success_since_startup_counter()
715                .get()
716        );
717
718        assert_eq!(
719            0 as f64,
720            metrics_service
721                .get_signer_registration_success_last_epoch_gauge()
722                .get()
723        );
724
725        assert_eq!(
726            1,
727            metrics_service.get_runtime_cycle_total_since_startup_counter().get()
728        );
729    }
730
731    #[tokio::test]
732    async fn registered_not_able_to_sign_to_unregistered() {
733        let mut runner = MockSignerRunner::new();
734        runner
735            .expect_get_current_time_point()
736            .once()
737            .returning(|| Ok(TimePoint::new(10, 100, ChainPoint::dummy())));
738        runner
739            .expect_update_era_checker()
740            .once()
741            .returning(|_e: Epoch| Ok(()));
742
743        let state_machine = init_state_machine(
744            SignerState::RegisteredNotAbleToSign { epoch: Epoch(0) },
745            runner,
746        );
747
748        state_machine
749            .cycle()
750            .await
751            .expect("Cycling the state machine should not fail");
752        assert_eq!(
753            SignerState::Unregistered { epoch: Epoch(10) },
754            state_machine.get_state().await
755        );
756    }
757
758    #[tokio::test]
759    async fn registered_not_able_to_sign_to_registered_not_able_to_sign() {
760        let mut runner = MockSignerRunner::new();
761        runner
762            .expect_get_current_time_point()
763            .once()
764            .returning(|| Ok(TimePoint::new(10, 100, ChainPoint::dummy())));
765
766        let state_machine = init_state_machine(
767            SignerState::RegisteredNotAbleToSign { epoch: Epoch(10) },
768            runner,
769        );
770
771        state_machine
772            .cycle()
773            .await
774            .expect("Cycling the state machine should not fail");
775        assert_eq!(
776            SignerState::RegisteredNotAbleToSign { epoch: Epoch(10) },
777            state_machine.get_state().await
778        );
779    }
780
781    #[tokio::test]
782    async fn ready_to_sign_to_unregistered() {
783        let mut runner = MockSignerRunner::new();
784        runner
785            .expect_get_current_time_point()
786            .once()
787            .returning(|| Ok(TimePoint::new(10, 100, ChainPoint::dummy())));
788        runner
789            .expect_update_era_checker()
790            .once()
791            .returning(|_e: Epoch| Ok(()));
792
793        let state_machine =
794            init_state_machine(SignerState::ReadyToSign { epoch: Epoch(0) }, runner);
795
796        state_machine
797            .cycle()
798            .await
799            .expect("Cycling the state machine should not fail");
800        assert_eq!(
801            SignerState::Unregistered { epoch: Epoch(10) },
802            state_machine.get_state().await
803        );
804    }
805
806    #[tokio::test]
807    async fn ready_to_sign_to_ready_to_sign_when_there_is_a_beacon_to_sign() {
808        let time_point = TimePoint::dummy();
809        let beacon_to_sign = BeaconToSign {
810            epoch: time_point.epoch,
811            signed_entity_type: SignedEntityType::MithrilStakeDistribution(time_point.epoch),
812            initiated_at: DateTime::default(),
813        };
814        let beacon_to_sign_clone = beacon_to_sign.clone();
815        let current_epoch = time_point.epoch;
816
817        let mut runner = MockSignerRunner::new();
818        runner
819            .expect_get_current_time_point()
820            .once()
821            .returning(move || Ok(time_point.to_owned()));
822        runner
823            .expect_get_beacon_to_sign()
824            .once()
825            .returning(move |_| Ok(Some(beacon_to_sign_clone.clone())));
826        runner
827            .expect_compute_message()
828            .once()
829            .returning(|_| Ok(ProtocolMessage::new()));
830        runner
831            .expect_compute_publish_single_signature()
832            .once()
833            .returning(|_, _| Ok(()));
834
835        let state_machine = init_state_machine(
836            SignerState::ReadyToSign {
837                epoch: current_epoch,
838            },
839            runner,
840        );
841        state_machine
842            .cycle()
843            .await
844            .expect("Cycling the state machine should not fail");
845
846        assert_eq!(
847            SignerState::ReadyToSign {
848                epoch: current_epoch
849            },
850            state_machine.get_state().await,
851            "state machine did not return a ReadyToSign state but {:?}",
852            state_machine.get_state().await
853        );
854    }
855
856    #[tokio::test]
857    async fn ready_to_sign_to_ready_to_sign_when_there_no_beacon_to_sign() {
858        let time_point = TimePoint::dummy();
859        let current_epoch = time_point.epoch;
860
861        let mut runner = MockSignerRunner::new();
862        runner
863            .expect_get_current_time_point()
864            .once()
865            .returning(move || Ok(time_point.to_owned()));
866        runner.expect_get_beacon_to_sign().once().returning(move |_| Ok(None));
867
868        let state_machine = init_state_machine(
869            SignerState::ReadyToSign {
870                epoch: current_epoch,
871            },
872            runner,
873        );
874        state_machine
875            .cycle()
876            .await
877            .expect("Cycling the state machine should not fail");
878
879        assert_eq!(
880            SignerState::ReadyToSign {
881                epoch: current_epoch
882            },
883            state_machine.get_state().await,
884            "state machine did not return a ReadyToSign state but {:?}",
885            state_machine.get_state().await
886        );
887    }
888}