zebra_network/protocol/external/
arbitrary.rs

Help
1//! Randomised test data generation for external protocol types.
2
3use std::net::SocketAddr;
4
5use proptest::{collection::vec, prelude::*};
6
7use zebra_chain::{block, transaction};
8
9use crate::PeerSocketAddr;
10
11use super::{
12    addr::canonical_peer_addr,
13    types::{PeerServices, Version},
14    InventoryHash, Message,
15};
16
17impl InventoryHash {
18    /// Generate a proptest strategy for [`InventoryHash::Error`]s.
19    pub fn error_strategy() -> BoxedStrategy<Self> {
20        Just(InventoryHash::Error).boxed()
21    }
22
23    /// Generate a proptest strategy for [`InventoryHash::Tx`] hashes.
24    pub fn tx_strategy() -> BoxedStrategy<Self> {
25        // using any::<transaction::Hash> causes a trait impl error
26        // when building the zebra-network crate separately
27        (any::<[u8; 32]>())
28            .prop_map(transaction::Hash)
29            .prop_map(InventoryHash::Tx)
30            .boxed()
31    }
32
33    /// Generate a proptest strategy for [`InventoryHash::Block`] hashes.
34    pub fn block_strategy() -> BoxedStrategy<Self> {
35        (any::<[u8; 32]>())
36            .prop_map(block::Hash)
37            .prop_map(InventoryHash::Block)
38            .boxed()
39    }
40
41    /// Generate a proptest strategy for [`InventoryHash::FilteredBlock`] hashes.
42    pub fn filtered_block_strategy() -> BoxedStrategy<Self> {
43        (any::<[u8; 32]>())
44            .prop_map(block::Hash)
45            .prop_map(InventoryHash::FilteredBlock)
46            .boxed()
47    }
48
49    /// Generate a proptest strategy for [`InventoryHash::Wtx`] hashes.
50    pub fn wtx_strategy() -> BoxedStrategy<Self> {
51        vec(any::<u8>(), 64)
52            .prop_map(|bytes| InventoryHash::Wtx(bytes.try_into().unwrap()))
53            .boxed()
54    }
55
56    /// Generate a proptest strategy for [`InventoryHash`] variants of the smallest serialized size.
57    pub fn smallest_types_strategy() -> BoxedStrategy<Self> {
58        InventoryHash::arbitrary()
59            .prop_filter(
60                "inventory type is not one of the smallest",
61                |inventory_hash| match inventory_hash {
62                    InventoryHash::Error
63                    | InventoryHash::Tx(_)
64                    | InventoryHash::Block(_)
65                    | InventoryHash::FilteredBlock(_) => true,
66                    InventoryHash::Wtx(_) => false,
67                },
68            )
69            .boxed()
70    }
71}
72
73impl Arbitrary for InventoryHash {
74    type Parameters = ();
75
76    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
77        prop_oneof![
78            Self::error_strategy(),
79            Self::tx_strategy(),
80            Self::block_strategy(),
81            Self::filtered_block_strategy(),
82            Self::wtx_strategy(),
83        ]
84        .boxed()
85    }
86
87    type Strategy = BoxedStrategy<Self>;
88}
89
90impl Arbitrary for PeerServices {
91    type Parameters = ();
92
93    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
94        any::<u64>()
95            .prop_map(PeerServices::from_bits_truncate)
96            .boxed()
97    }
98
99    type Strategy = BoxedStrategy<Self>;
100}
101
102impl Message {
103    /// Create a strategy that only generates [`Message::Inv`] messages.
104    pub fn inv_strategy() -> BoxedStrategy<Self> {
105        any::<Vec<InventoryHash>>().prop_map(Message::Inv).boxed()
106    }
107
108    /// Create a strategy that only generates [`Message::GetData`] messages.
109    pub fn get_data_strategy() -> BoxedStrategy<Self> {
110        any::<Vec<InventoryHash>>()
111            .prop_map(Message::GetData)
112            .boxed()
113    }
114}
115
116impl Arbitrary for Version {
117    type Parameters = ();
118
119    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
120        prop_oneof![170_002_u32..=170_100, 0_u32..]
121            .prop_map(Version)
122            .boxed()
123    }
124
125    type Strategy = BoxedStrategy<Self>;
126}
127
128/// Returns a random canonical Zebra [`PeerSocketAddr`].
129///
130/// See [`canonical_ip_addr`] for details.
131///
132/// [`canonical_ip_addr`]: super::addr::canonical::canonical_ip_addr
133pub fn canonical_peer_addr_strategy() -> impl Strategy<Value = PeerSocketAddr> {
134    any::<SocketAddr>().prop_map(canonical_peer_addr)
135}
zebra_network/protocol/external/arbitrary.rs

zebra_network/protocol/external/
arbitrary.rs
