diff --git a/node/actors/bft/src/testonly/ut_harness.rs b/node/actors/bft/src/testonly/ut_harness.rs index e550e8db..1ab48266 100644 --- a/node/actors/bft/src/testonly/ut_harness.rs +++ b/node/actors/bft/src/testonly/ut_harness.rs @@ -5,7 +5,11 @@ use crate::{ Consensus, }; use rand::{rngs::StdRng, Rng}; -use zksync_concurrency::{ctx, ctx::Ctx, scope}; +use zksync_concurrency::{ + ctx, + ctx::{Canceled, Ctx}, + scope, +}; use zksync_consensus_network::io::ConsensusInputMessage; use zksync_consensus_roles::validator::{ self, BlockHeader, CommitQC, ConsensusMsg, LeaderCommit, LeaderPrepare, Payload, Phase, @@ -273,13 +277,12 @@ impl UTHarness { .unwrap() } - pub(crate) async fn recv_signed(&mut self) -> Option> { - let msg = self.pipe.recv(&self.ctx).await.unwrap(); - match msg { + pub(crate) async fn recv_signed(&mut self) -> Result, Canceled> { + self.pipe.recv(&self.ctx).await.map(|output_message| match output_message { OutputMessage::Network(ConsensusInputMessage { message: signed, .. - }) => Some(signed), - } + }) => signed, + }) } pub(crate) fn current_replica_view(&self) -> ViewNumber {