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

A method that returns a proptest strategy that creates
`Message::GetData` instances.
---
 zebra-network/src/protocol/external/arbitrary.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/zebra-network/src/protocol/external/arbitrary.rs b/zebra-network/src/protocol/external/arbitrary.rs
index a0ebe48d25e..4a20e34ac1f 100644
--- a/zebra-network/src/protocol/external/arbitrary.rs
+++ b/zebra-network/src/protocol/external/arbitrary.rs
@@ -96,4 +96,11 @@ impl Message {
     pub fn inv_strategy() -> BoxedStrategy<Self> {
         any::<Vec<InventoryHash>>().prop_map(Message::Inv).boxed()
     }
+
+    /// Create a strategy that only generates [`Message::GetData`] messages.
+    pub fn get_data_strategy() -> BoxedStrategy<Self> {
+        any::<Vec<InventoryHash>>()
+            .prop_map(Message::GetData)
+            .boxed()
+    }
 }