test(consensus): update sm tests

starkware-libs · Jul 25, 2024 · 3b543c2 · 3b543c2
1 parent e80cd8e
commit 3b543c2
Showing 1 changed file with 17 additions and 33 deletions.
diff --git a/crates/sequencing/papyrus_consensus/src/state_machine_test.rs b/crates/sequencing/papyrus_consensus/src/state_machine_test.rs
@@ -8,7 +8,8 @@ use crate::state_machine::{StateMachine, StateMachineEvent};
 use crate::types::ValidatorId;
 
 lazy_static! {
-    static ref VALIDATOR_ID: ValidatorId = 0_u32.into();
+    static ref PROPOSER_ID: ValidatorId = 0_u32.into();
+    static ref VALIDATOR_ID: ValidatorId = 1_u32.into();
 }
 
 const BLOCK_HASH: Option<BlockHash> = Some(BlockHash(Felt::ONE));
@@ -17,17 +18,17 @@ const ROUND: Round = 0;
 #[test_case(true; "proposer")]
 #[test_case(false; "validator")]
 fn events_arrive_in_ideal_order(is_proposer: bool) {
-    let mut state_machine = StateMachine::new(*VALIDATOR_ID, 4);
-    let leader_fn = |_: Round| *VALIDATOR_ID;
+    let id = if is_proposer { *PROPOSER_ID } else { *VALIDATOR_ID };
+    let mut state_machine = StateMachine::new(id, 4);
+    let leader_fn = |_: Round| *PROPOSER_ID;
     let mut events = state_machine.start(&leader_fn);
-    assert_eq!(events.pop_front().unwrap(), StateMachineEvent::GetProposal(None, ROUND));
     if is_proposer {
+        assert_eq!(events.pop_front().unwrap(), StateMachineEvent::GetProposal(None, ROUND));
         events = state_machine
             .handle_event(StateMachineEvent::GetProposal(BLOCK_HASH, ROUND), &leader_fn);
         assert_eq!(events.pop_front().unwrap(), StateMachineEvent::Proposal(BLOCK_HASH, ROUND));
     } else {
-        state_machine.handle_event(StateMachineEvent::GetProposal(None, ROUND), &leader_fn);
-        assert!(events.is_empty());
+        assert!(events.is_empty(), "{:?}", events);
         events =
             state_machine.handle_event(StateMachineEvent::Proposal(BLOCK_HASH, ROUND), &leader_fn);
     }
@@ -57,18 +58,12 @@ fn events_arrive_in_ideal_order(is_proposer: bool) {
 #[test]
 fn validator_receives_votes_first() {
     let mut state_machine = StateMachine::new(*VALIDATOR_ID, 4);
-
-    let leader_fn = |_: Round| *VALIDATOR_ID;
-    let mut events = state_machine.start(&leader_fn);
-    assert_eq!(events.pop_front().unwrap(), StateMachineEvent::GetProposal(None, ROUND));
-    assert!(events.is_empty(), "{:?}", events);
-    events = state_machine.handle_event(StateMachineEvent::GetProposal(None, ROUND), &leader_fn);
-    assert!(events.is_empty(), "{:?}", events);
+    let leader_fn = |_: Round| *PROPOSER_ID;
 
     // Receives votes from all the other nodes first (more than minimum for a quorum).
-    events.append(
-        &mut state_machine.handle_event(StateMachineEvent::Prevote(BLOCK_HASH, ROUND), &leader_fn),
-    );
+    let mut events =
+        state_machine.handle_event(StateMachineEvent::Prevote(BLOCK_HASH, ROUND), &leader_fn);
+
     events.append(
         &mut state_machine.handle_event(StateMachineEvent::Prevote(BLOCK_HASH, ROUND), &leader_fn),
     );
@@ -101,32 +96,21 @@ fn validator_receives_votes_first() {
 }
 
 #[test]
-fn buffer_events_during_get_proposal() {
+fn buffer_events_during_proposal() {
     let mut state_machine = StateMachine::new(*VALIDATOR_ID, 4);
-    let leader_fn = |_: Round| *VALIDATOR_ID;
-    let mut events = state_machine.start(&leader_fn);
-    assert_eq!(events.pop_front().unwrap(), StateMachineEvent::GetProposal(None, 0));
+    let leader_fn = |_: Round| *PROPOSER_ID;
+    let mut events =
+        state_machine.handle_event(StateMachineEvent::Proposal(BLOCK_HASH, ROUND), &leader_fn);
+    assert_eq!(events.pop_front().unwrap(), StateMachineEvent::Prevote(BLOCK_HASH, ROUND));
     assert!(events.is_empty(), "{:?}", events);
 
-    // TODO(matan): When we support NIL votes, we should send them. Real votes without the proposal
-    // doesn't make sense.
-    events.append(
-        &mut state_machine.handle_event(StateMachineEvent::Proposal(BLOCK_HASH, ROUND), &leader_fn),
-    );
+    // TODO(matan): When we support NIL votes, we should send them.
     events.append(
         &mut state_machine.handle_event(StateMachineEvent::Prevote(BLOCK_HASH, ROUND), &leader_fn),
     );
     events.append(
         &mut state_machine.handle_event(StateMachineEvent::Prevote(BLOCK_HASH, ROUND), &leader_fn),
     );
-    events.append(
-        &mut state_machine.handle_event(StateMachineEvent::Prevote(BLOCK_HASH, ROUND), &leader_fn),
-    );
-    assert!(events.is_empty(), "{:?}", events);
-
-    // Node finishes building the proposal.
-    events = state_machine.handle_event(StateMachineEvent::GetProposal(None, 0), &leader_fn);
-    assert_eq!(events.pop_front().unwrap(), StateMachineEvent::Prevote(BLOCK_HASH, ROUND));
     assert_eq!(events.pop_front().unwrap(), StateMachineEvent::Precommit(BLOCK_HASH, ROUND));
     assert!(events.is_empty(), "{:?}", events);
 }