chore(sequencing): remove mock context from manager_test.rs

crates/sequencing/papyrus_consensus/src/manager_test.rs

@@ -1,40 +1,25 @@
 use std::time::Duration;
 use std::vec;
 
-use async_trait::async_trait;
 use futures::channel::{mpsc, oneshot};
 use futures::SinkExt;
 use lazy_static::lazy_static;
-use mockall::mock;
 use mockall::predicate::eq;
 use papyrus_network::network_manager::test_utils::{
     mock_register_broadcast_topic,
     MockBroadcastedMessagesSender,
     TestSubscriberChannels,
 };
 use papyrus_network_types::network_types::BroadcastedMessageMetadata;
-use papyrus_protobuf::consensus::{
-    ConsensusMessage,
-    ProposalFin,
-    ProposalInit,
-    ProposalPart,
-    Vote,
-};
+use papyrus_protobuf::consensus::{ConsensusMessage, ProposalFin};
 use papyrus_test_utils::{get_rng, GetTestInstance};
 use starknet_api::block::{BlockHash, BlockNumber};
 use starknet_types_core::felt::Felt;
 
 use super::{run_consensus, MultiHeightManager};
 use crate::config::TimeoutsConfig;
-use crate::test_utils::{precommit, prevote, proposal_init};
-use crate::types::{
-    ConsensusContext,
-    ConsensusError,
-    ProposalContentId,
-    Round,
-    ValidatorId,
-    DEFAULT_VALIDATOR_ID,
-};
+use crate::test_utils::{precommit, prevote, proposal_init, MockTestContext, TestProposalPart};
+use crate::types::{ConsensusError, ValidatorId, DEFAULT_VALIDATOR_ID};
 
 lazy_static! {
     static ref PROPOSER_ID: ValidatorId = DEFAULT_VALIDATOR_ID.into();
@@ -46,59 +31,15 @@ lazy_static! {
 
 const CHANNEL_SIZE: usize = 10;
 
-mock! {
-    pub TestContext {}
-
-    #[async_trait]
-    impl ConsensusContext for TestContext {
-        type ProposalPart = ProposalPart;
-
-        async fn build_proposal(
-            &mut self,
-            init: ProposalInit,
-            timeout: Duration
-        ) -> oneshot::Receiver<ProposalContentId>;
-
-        async fn validate_proposal(
-            &mut self,
-            height: BlockNumber,
-            round: Round,
-            proposer: ValidatorId,
-            timeout: Duration,
-            content: mpsc::Receiver<ProposalPart>
-        ) -> oneshot::Receiver<(ProposalContentId, ProposalFin)>;
-
-        async fn repropose(
-            &mut self,
-            id: ProposalContentId,
-            init: ProposalInit,
-        );
-
-        async fn validators(&self, height: BlockNumber) -> Vec<ValidatorId>;
-
-        fn proposer(&self, height: BlockNumber, round: Round) -> ValidatorId;
-
-        async fn broadcast(&mut self, message: ConsensusMessage) -> Result<(), ConsensusError>;
-
-        async fn decision_reached(
-            &mut self,
-            block: ProposalContentId,
-            precommits: Vec<Vote>,
-        ) -> Result<(), ConsensusError>;
-
-        async fn set_height_and_round(&mut self, height: BlockNumber, round: Round);
-    }
-}
-
 async fn send(sender: &mut MockBroadcastedMessagesSender<ConsensusMessage>, msg: ConsensusMessage) {
     let broadcasted_message_metadata =
         BroadcastedMessageMetadata::get_test_instance(&mut get_rng());
     sender.send((msg, broadcasted_message_metadata)).await.unwrap();
 }
 
 async fn send_proposal(
-    proposal_receiver_sender: &mut mpsc::Sender<mpsc::Receiver<ProposalPart>>,
-    content: Vec<ProposalPart>,
+    proposal_receiver_sender: &mut mpsc::Sender<mpsc::Receiver<TestProposalPart>>,
+    content: Vec<TestProposalPart>,
 ) {
     let (mut proposal_sender, proposal_receiver) = mpsc::channel(CHANNEL_SIZE);
     proposal_receiver_sender.send(proposal_receiver).await.unwrap();
@@ -137,8 +78,8 @@ async fn manager_multiple_heights_unordered() {
     send_proposal(
         &mut proposal_receiver_sender,
         vec![
-            ProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID)),
-            ProposalPart::Fin(ProposalFin { proposal_content_id: BlockHash(Felt::TWO) }),
+            TestProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID)),
+            TestProposalPart::Fin(ProposalFin { proposal_content_id: BlockHash(Felt::TWO) }),
         ],
     )
     .await;
@@ -148,8 +89,8 @@ async fn manager_multiple_heights_unordered() {
     send_proposal(
         &mut proposal_receiver_sender,
         vec![
-            ProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID)),
-            ProposalPart::Fin(ProposalFin { proposal_content_id: BlockHash(Felt::ONE) }),
+            TestProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID)),
+            TestProposalPart::Fin(ProposalFin { proposal_content_id: BlockHash(Felt::ONE) }),
         ],
     )
     .await;
@@ -217,7 +158,7 @@ async fn run_consensus_sync() {
     // Send messages for height 2.
     send_proposal(
         &mut proposal_receiver_sender,
-        vec![ProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID))],
+        vec![TestProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID))],
     )
     .await;
     let TestSubscriberChannels { mock_network, subscriber_channels } =
@@ -308,7 +249,7 @@ async fn run_consensus_sync_cancellation_safety() {
     // Send a proposal for height 1.
     send_proposal(
         &mut proposal_receiver_sender,
-        vec![ProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID))],
+        vec![TestProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID))],
     )
     .await;
     proposal_handled_rx.await.unwrap();
@@ -340,7 +281,7 @@ async fn test_timeouts() {
 
     send_proposal(
         &mut proposal_receiver_sender,
-        vec![ProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID))],
+        vec![TestProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID))],
     )
     .await;
     send(&mut sender, prevote(None, 1, 0, *VALIDATOR_ID_2)).await;
@@ -395,7 +336,7 @@ async fn test_timeouts() {
     // reach a decision.
     send_proposal(
         &mut proposal_receiver_sender,
-        vec![ProposalPart::Init(proposal_init(1, 1, *PROPOSER_ID))],
+        vec![TestProposalPart::Init(proposal_init(1, 1, *PROPOSER_ID))],
     )
     .await;
     send(&mut sender, prevote(Some(Felt::ONE), 1, 1, *PROPOSER_ID)).await;

crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs

@@ -11,7 +11,7 @@ use super::SingleHeightConsensus;
 use crate::config::TimeoutsConfig;
 use crate::single_height_consensus::{ShcEvent, ShcReturn, ShcTask};
 use crate::state_machine::StateMachineEvent;
-use crate::test_utils::{precommit, prevote, MockProposalPart, MockTestContext, TestBlock};
+use crate::test_utils::{precommit, prevote, MockTestContext, TestBlock, TestProposalPart};
 use crate::types::{ConsensusError, ValidatorId, DEFAULT_VALIDATOR_ID};
 
 lazy_static! {
@@ -69,7 +69,7 @@ async fn handle_proposal(
 ) -> ShcReturn {
     // Send the proposal from the peer.
     let (mut content_sender, content_receiver) = mpsc::channel(CHANNEL_SIZE);
-    content_sender.send(MockProposalPart(1)).await.unwrap();
+    content_sender.send(TestProposalPart::Init(ProposalInit::default())).await.unwrap();
 
     shc.handle_proposal(context, PROPOSAL_INIT.clone(), content_receiver).await.unwrap()
 }

crates/sequencing/papyrus_consensus/src/test_utils.rs

@@ -8,14 +8,7 @@ use papyrus_protobuf::converters::ProtobufConversionError;
 use starknet_api::block::{BlockHash, BlockNumber};
 use starknet_types_core::felt::Felt;
 
-use crate::types::{
-    ConsensusContext,
-    ConsensusError,
-    ProposalContentId,
-    Round,
-    ValidatorId,
-    DEFAULT_VALIDATOR_ID,
-};
+use crate::types::{ConsensusContext, ConsensusError, ProposalContentId, Round, ValidatorId};
 
 /// Define a consensus block which can be used to enable auto mocking Context.
 #[derive(Debug, PartialEq, Clone)]
@@ -25,36 +18,46 @@ pub struct TestBlock {
 }
 
 #[derive(Debug, PartialEq, Clone)]
-pub struct MockProposalPart(pub u64);
+pub enum TestProposalPart {
+    Init(ProposalInit),
+    Fin(ProposalFin),
+}
 
-impl From<ProposalInit> for MockProposalPart {
+impl From<ProposalInit> for TestProposalPart {
     fn from(init: ProposalInit) -> Self {
-        MockProposalPart(init.height.0)
+        TestProposalPart::Init(init)
     }
 }
 
-impl TryFrom<MockProposalPart> for ProposalInit {
+impl TryFrom<TestProposalPart> for ProposalInit {
     type Error = ProtobufConversionError;
-    fn try_from(part: MockProposalPart) -> Result<Self, Self::Error> {
-        Ok(ProposalInit {
-            height: BlockNumber(part.0),
-            proposer: DEFAULT_VALIDATOR_ID.into(),
-            ..Default::default()
-        })
+    fn try_from(part: TestProposalPart) -> Result<Self, Self::Error> {
+        match part {
+            TestProposalPart::Init(init) => Ok(init),
+            _ => Err(ProtobufConversionError::WrongEnumVariant {
+                type_description: "TestProposalPart",
+                expected: "Init",
+                value_as_str: format!("{:?}", part),
+            }),
+        }
     }
 }
 
-impl From<MockProposalPart> for Vec<u8> {
-    fn from(part: MockProposalPart) -> Vec<u8> {
-        vec![u8::try_from(part.0).expect("Invalid MockProposalPart conversion")]
+impl From<TestProposalPart> for Vec<u8> {
+    fn from(part: TestProposalPart) -> Vec<u8> {
+        let init = match part {
+            TestProposalPart::Init(init) => init,
+            _ => panic!("Invalid TestProposalPart conversion"),
+        };
+        <Vec<u8>>::try_from(init).expect("Invalid TestProposalPart conversion")
     }
 }
 
-impl TryFrom<Vec<u8>> for MockProposalPart {
+impl TryFrom<Vec<u8>> for TestProposalPart {
     type Error = ProtobufConversionError;
 
     fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
-        Ok(MockProposalPart(value[0].into()))
+        Ok(TestProposalPart::Init(value.try_into()?))
     }
 }
 
@@ -64,7 +67,7 @@ mock! {
 
     #[async_trait]
     impl ConsensusContext for TestContext {
-        type ProposalPart = MockProposalPart;
+        type ProposalPart = TestProposalPart;
 
         async fn build_proposal(
             &mut self,
@@ -78,7 +81,7 @@ mock! {
             round: Round,
             proposer: ValidatorId,
             timeout: Duration,
-            content: mpsc::Receiver<MockProposalPart>
+            content: mpsc::Receiver<TestProposalPart>
         ) -> oneshot::Receiver<(ProposalContentId, ProposalFin)>;
 
         async fn repropose(