Skip to content

Commit

Permalink
chore(sequencing): remove ConsensusMessage from stream_handler_test
Browse files Browse the repository at this point in the history
  • Loading branch information
guy-starkware committed Dec 10, 2024
1 parent 06a3192 commit 0a4ca1d
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 20 deletions.
1 change: 1 addition & 0 deletions crates/papyrus_protobuf/src/converters/consensus_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ fn convert_stream_message_to_vec_u8_and_back() {
assert_eq!(stream_message, res_data);
}

// TODO(guyn): this can be removed once ConsensusMessage is taken out.
#[test]
fn convert_consensus_message_to_vec_u8_and_back() {
let mut rng = get_rng();
Expand Down
5 changes: 3 additions & 2 deletions crates/papyrus_protobuf/src/converters/test_instances.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ use starknet_api::core::ContractAddress;
use starknet_api::transaction::{Transaction, TransactionHash};

use crate::consensus::{
ConsensusMessage,
Proposal,
ConsensusMessage, // TODO: remove this
Proposal, // TODO: remove this
ProposalFin,
ProposalInit,
ProposalPart,
Expand All @@ -18,6 +18,7 @@ use crate::consensus::{
};

auto_impl_get_test_instance! {
// TODO(guyn): remove this once we integrate ProposalPart everywhere.
pub enum ConsensusMessage {
Proposal(Proposal) = 0,
Vote(Vote) = 1,
Expand Down
35 changes: 17 additions & 18 deletions crates/sequencing/papyrus_consensus/src/stream_handler_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use papyrus_network::network_manager::test_utils::{
};
use papyrus_network::network_manager::BroadcastTopicChannels;
use papyrus_network_types::network_types::BroadcastedMessageMetadata;
use papyrus_protobuf::consensus::{ConsensusMessage, Proposal, StreamMessage, StreamMessageBody};
use papyrus_protobuf::consensus::{StreamMessage, StreamMessageBody};
use papyrus_test_utils::{get_rng, GetTestInstance};

use super::{MessageId, StreamHandler, StreamId};
Expand All @@ -20,16 +20,18 @@ const CHANNEL_SIZE: usize = 100;

#[cfg(test)]
mod tests {
use papyrus_protobuf::consensus::{ProposalInit, ProposalPart};

use super::*;

fn make_test_message(
stream_id: StreamId,
message_id: MessageId,
fin: bool,
) -> StreamMessage<ConsensusMessage> {
) -> StreamMessage<ProposalPart> {
let content = match fin {
true => StreamMessageBody::Fin,
false => StreamMessageBody::Content(ConsensusMessage::Proposal(Proposal::default())),
false => StreamMessageBody::Content(ProposalPart::Init(ProposalInit::default())),
};
StreamMessage { message: content, stream_id, message_id }
}
Expand All @@ -48,24 +50,21 @@ mod tests {
}

async fn send(
sender: &mut MockBroadcastedMessagesSender<StreamMessage<ConsensusMessage>>,
sender: &mut MockBroadcastedMessagesSender<StreamMessage<ProposalPart>>,
metadata: &BroadcastedMessageMetadata,
msg: StreamMessage<ConsensusMessage>,
msg: StreamMessage<ProposalPart>,
) {
sender.send((msg, metadata.clone())).await.unwrap();
}

#[allow(clippy::type_complexity)]
fn setup_test() -> (
StreamHandler<ConsensusMessage>,
MockBroadcastedMessagesSender<StreamMessage<ConsensusMessage>>,
mpsc::Receiver<mpsc::Receiver<ConsensusMessage>>,
StreamHandler<ProposalPart>,
MockBroadcastedMessagesSender<StreamMessage<ProposalPart>>,
mpsc::Receiver<mpsc::Receiver<ProposalPart>>,
BroadcastedMessageMetadata,
mpsc::Sender<(StreamId, mpsc::Receiver<ConsensusMessage>)>,
futures::stream::Map<
mpsc::Receiver<Vec<u8>>,
fn(Vec<u8>) -> StreamMessage<ConsensusMessage>,
>,
mpsc::Sender<(StreamId, mpsc::Receiver<ProposalPart>)>,
futures::stream::Map<mpsc::Receiver<Vec<u8>>, fn(Vec<u8>) -> StreamMessage<ProposalPart>>,
) {
// The outbound_sender is the network connector for broadcasting messages.
// The network_broadcast_receiver is used to catch those messages in the test.
Expand All @@ -82,7 +81,7 @@ mod tests {
// The receiver goes into StreamHandler, sender is used by the test (as mock Consensus).
// Note that each new channel comes in a tuple with (stream_id, receiver).
let (outbound_channel_sender, outbound_channel_receiver) =
mpsc::channel::<(StreamId, mpsc::Receiver<ConsensusMessage>)>(CHANNEL_SIZE);
mpsc::channel::<(StreamId, mpsc::Receiver<ProposalPart>)>(CHANNEL_SIZE);

// The network_sender_to_inbound is the sender of the mock network, that is used by the
// test to send messages into the StreamHandler (from the mock network).
Expand All @@ -100,7 +99,7 @@ mod tests {
// each stream. The inbound_channel_receiver is given to the "mock consensus" that
// gets new channels and inbounds to them.
let (inbound_channel_sender, inbound_channel_receiver) =
mpsc::channel::<mpsc::Receiver<ConsensusMessage>>(CHANNEL_SIZE);
mpsc::channel::<mpsc::Receiver<ProposalPart>>(CHANNEL_SIZE);

// TODO(guyn): We should also give the broadcast_topic_client to the StreamHandler
// This will allow reporting to the network things like bad peers.
Expand Down Expand Up @@ -419,7 +418,7 @@ mod tests {
broadcast_channel_sender.send((stream_id1, receiver1)).await.unwrap();

// Send a message on the stream.
let message1 = ConsensusMessage::Proposal(Proposal::default());
let message1 = ProposalPart::Init(ProposalInit::default());
sender1.send(message1.clone()).await.unwrap();

// Run the loop for a short duration to process the message.
Expand All @@ -446,7 +445,7 @@ mod tests {
assert_eq!(stream_handler.outbound_stream_number[&stream_id1], 1);

// Send another message on the same stream.
let message2 = ConsensusMessage::Proposal(Proposal::default());
let message2 = ProposalPart::Init(ProposalInit::default());
sender1.send(message2.clone()).await.unwrap();

// Run the loop for a short duration to process the message.
Expand All @@ -471,7 +470,7 @@ mod tests {
broadcast_channel_sender.send((stream_id2, receiver2)).await.unwrap();

// Send a message on the stream.
let message3 = ConsensusMessage::Proposal(Proposal::default());
let message3 = ProposalPart::Init(ProposalInit::default());
sender2.send(message3.clone()).await.unwrap();

// Run the loop for a short duration to process the message.
Expand Down

0 comments on commit 0a4ca1d

Please sign in to comment.