diff --git a/zebra-network/src/protocol/external/arbitrary.rs b/zebra-network/src/protocol/external/arbitrary.rs index ceb9edeb1e4..a0ebe48d25e 100644 --- a/zebra-network/src/protocol/external/arbitrary.rs +++ b/zebra-network/src/protocol/external/arbitrary.rs @@ -2,7 +2,7 @@ use std::convert::TryInto; use proptest::{arbitrary::any, arbitrary::Arbitrary, collection::vec, prelude::*}; -use super::{types::PeerServices, InventoryHash}; +use super::{types::PeerServices, InventoryHash, Message}; use zebra_chain::{block, transaction}; @@ -90,3 +90,10 @@ impl Arbitrary for PeerServices { type Strategy = BoxedStrategy; } + +impl Message { + /// Create a strategy that only generates [`Message::Inv`] messages. + pub fn inv_strategy() -> BoxedStrategy { + any::>().prop_map(Message::Inv).boxed() + } +}