From 3f1ba69e05bd2b5ca36fddd62bae8975cba537dd Mon Sep 17 00:00:00 2001
From: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>
Date: Mon, 5 Jul 2021 20:22:34 +0000
Subject: [PATCH] Implement `Message::inv_strategy()`

A method to return a proptest strategy that creates `Message::Inv`
instances.
---
 zebra-network/src/protocol/external/arbitrary.rs | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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<Self>;
 }
+
+impl Message {
+    /// Create a strategy that only generates [`Message::Inv`] messages.
+    pub fn inv_strategy() -> BoxedStrategy<Self> {
+        any::<Vec<InventoryHash>>().prop_map(Message::Inv).boxed()
+    }
+}