diff --git a/crates/sequencing/papyrus_consensus/src/manager_test.rs b/crates/sequencing/papyrus_consensus/src/manager_test.rs index e6ae664568..b9de4dbb35 100644 --- a/crates/sequencing/papyrus_consensus/src/manager_test.rs +++ b/crates/sequencing/papyrus_consensus/src/manager_test.rs @@ -2,11 +2,9 @@ use std::sync::Arc; 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, @@ -14,13 +12,7 @@ use papyrus_network::network_manager::test_utils::{ 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; @@ -28,15 +20,8 @@ use tokio::sync::Notify; use super::{run_consensus, MultiHeightManager, RunHeightRes}; 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(); @@ -52,50 +37,6 @@ 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; - - async fn validate_proposal( - &mut self, - height: BlockNumber, - round: Round, - proposer: ValidatorId, - timeout: Duration, - content: mpsc::Receiver - ) -> oneshot::Receiver<(ProposalContentId, ProposalFin)>; - - async fn repropose( - &mut self, - id: ProposalContentId, - init: ProposalInit, - ); - - async fn validators(&self, height: BlockNumber) -> Vec; - - 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, - ) -> Result<(), ConsensusError>; - - async fn set_height_and_round(&mut self, height: BlockNumber, round: Round); - } -} - async fn send(sender: &mut MockBroadcastedMessagesSender, msg: ConsensusMessage) { let broadcasted_message_metadata = BroadcastedMessageMetadata::get_test_instance(&mut get_rng()); @@ -103,8 +44,8 @@ async fn send(sender: &mut MockBroadcastedMessagesSender, msg: } async fn send_proposal( - proposal_receiver_sender: &mut mpsc::Sender>, - content: Vec, + proposal_receiver_sender: &mut mpsc::Sender>, + content: Vec, ) { let (mut proposal_sender, proposal_receiver) = mpsc::channel(CHANNEL_SIZE); proposal_receiver_sender.send(proposal_receiver).await.unwrap(); @@ -149,10 +90,7 @@ async fn manager_multiple_heights_unordered() { // Send messages for height 2 followed by those for height 1. send_proposal( &mut proposal_receiver_sender, - vec![ - ProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID)), - ProposalPart::Fin(ProposalFin { proposal_content_id: BlockHash(Felt::TWO) }), - ], + vec![TestProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID))], ) .await; send(&mut sender, prevote(Some(Felt::TWO), 2, 0, *PROPOSER_ID)).await; @@ -160,10 +98,7 @@ 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) }), - ], + vec![TestProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID))], ) .await; send(&mut sender, prevote(Some(Felt::ONE), 1, 0, *PROPOSER_ID)).await; @@ -232,7 +167,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 } = @@ -326,7 +261,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.notified().await; @@ -358,7 +293,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; @@ -414,7 +349,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; diff --git a/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs b/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs index ee5f2cd062..8f575fe6d7 100644 --- a/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs +++ b/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() } diff --git a/crates/sequencing/papyrus_consensus/src/test_utils.rs b/crates/sequencing/papyrus_consensus/src/test_utils.rs index 8ee68c1427..f1f265bf32 100644 --- a/crates/sequencing/papyrus_consensus/src/test_utils.rs +++ b/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,36 @@ pub struct TestBlock { } #[derive(Debug, PartialEq, Clone)] -pub struct MockProposalPart(pub u64); +pub enum TestProposalPart { + Init(ProposalInit), +} -impl From for MockProposalPart { +impl From for TestProposalPart { fn from(init: ProposalInit) -> Self { - MockProposalPart(init.height.0) + TestProposalPart::Init(init) } } -impl TryFrom for ProposalInit { +impl TryFrom for ProposalInit { type Error = ProtobufConversionError; - fn try_from(part: MockProposalPart) -> Result { - Ok(ProposalInit { - height: BlockNumber(part.0), - proposer: DEFAULT_VALIDATOR_ID.into(), - ..Default::default() - }) + fn try_from(part: TestProposalPart) -> Result { + let TestProposalPart::Init(init) = part; + Ok(init) } } -impl From for Vec { - fn from(part: MockProposalPart) -> Vec { - vec![u8::try_from(part.0).expect("Invalid MockProposalPart conversion")] +impl From for Vec { + fn from(part: TestProposalPart) -> Vec { + let TestProposalPart::Init(init) = part; + init.try_into().expect("Invalid TestProposalPart conversion") } } -impl TryFrom> for MockProposalPart { +impl TryFrom> for TestProposalPart { type Error = ProtobufConversionError; fn try_from(value: Vec) -> Result { - Ok(MockProposalPart(value[0].into())) + Ok(TestProposalPart::Init(value.try_into()?)) } } @@ -64,7 +57,7 @@ mock! { #[async_trait] impl ConsensusContext for TestContext { - type ProposalPart = MockProposalPart; + type ProposalPart = TestProposalPart; async fn build_proposal( &mut self, @@ -78,7 +71,7 @@ mock! { round: Round, proposer: ValidatorId, timeout: Duration, - content: mpsc::Receiver + content: mpsc::Receiver ) -> oneshot::Receiver<(ProposalContentId, ProposalFin)>; async fn repropose(