Skip to content

Commit

Permalink
chore(sequencing): remove mock context from manager_test.rs (#2524)
Browse files Browse the repository at this point in the history 
I've removed the MockTestContext from manager_test.rs and instead moved it to depend on the mock context in test_utils.rs.
  • Loading branch information
@guy-starkware
guy-starkware authored Dec 20, 2024
1 parent 1dfd8a7 commit c3325e5
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 83 deletions.
79 changes: 11 additions & 68 deletions crates/sequencing/papyrus_consensus/src/manager_test.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,35 +2,26 @@ 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,
MockBroadcastedMessagesSender,
TestSubscriberChannels,
};
use papyrus_network_types::network_types::BroadcastedMessageMetadata;
use papyrus_protobuf::consensus::{
ConsensusMessage,
ProposalFin,
ProposalInit,
ProposalPart,
Vote,
DEFAULT_VALIDATOR_ID,
};
use papyrus_protobuf::consensus::{ConsensusMessage, ProposalFin, DEFAULT_VALIDATOR_ID};
use papyrus_test_utils::{get_rng, GetTestInstance};
use starknet_api::block::{BlockHash, BlockNumber};
use starknet_types_core::felt::Felt;
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};
use crate::test_utils::{precommit, prevote, proposal_init, MockTestContext, TestProposalPart};
use crate::types::{ConsensusError, ValidatorId};

lazy_static! {
static ref PROPOSER_ID: ValidatorId = DEFAULT_VALIDATOR_ID.into();
Expand All @@ -46,57 +37,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,
init: ProposalInit,
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();
Expand Down Expand Up @@ -141,21 +90,15 @@ 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;
send(&mut sender, precommit(Some(Felt::TWO), 2, 0, *PROPOSER_ID)).await;

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;
Expand Down Expand Up @@ -224,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 } =
Expand Down Expand Up @@ -318,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;
Expand Down Expand Up @@ -350,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;
Expand Down Expand Up @@ -406,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;
Expand Down
4 changes: 2 additions & 2 deletions crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,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};

lazy_static! {
Expand Down Expand Up @@ -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, content_receiver).await.unwrap()
}
Expand Down
30 changes: 17 additions & 13 deletions crates/sequencing/papyrus_consensus/src/test_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,32 +18,36 @@ pub struct TestBlock {
}

#[derive(Debug, PartialEq, Clone)]
pub struct MockProposalPart(pub u64);
pub enum TestProposalPart {
Init(ProposalInit),
}

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), ..Default::default() })
fn try_from(part: TestProposalPart) -> Result<Self, Self::Error> {
let TestProposalPart::Init(init) = part;
Ok(init)
}
}

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 TestProposalPart::Init(init) = part;
init.into()
}
}

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()?))
}
}

Expand All @@ -53,7 +57,7 @@ mock! {

#[async_trait]
impl ConsensusContext for TestContext {
type ProposalPart = MockProposalPart;
type ProposalPart = TestProposalPart;

async fn build_proposal(
&mut self,
Expand All @@ -65,7 +69,7 @@ mock! {
&mut self,
init: ProposalInit,
timeout: Duration,
content: mpsc::Receiver<MockProposalPart>
content: mpsc::Receiver<TestProposalPart>
) -> oneshot::Receiver<(ProposalContentId, ProposalFin)>;

async fn repropose(
Expand Down

0 comments on commit c3325e5

Please sign in to comment.