Skip to main content

spate_test/
strategies.rs

1//! Property-testing strategies for framework types (feature `proptest`).
2//!
3//! Use these to drive the mocks from generated inputs, e.g. pushing
4//! arbitrary payload sets through a pipeline and asserting invariants like
5//! "every acknowledged offset was durably written".
6
7use crate::sink::WriteOutcome;
8use proptest::prelude::*;
9use spate_core::record::PartitionId;
10use spate_core::source::LaneId;
11use std::time::Duration;
12
13/// Up to `max_items` payloads of up to `max_len` bytes each.
14pub fn payloads(max_items: usize, max_len: usize) -> impl Strategy<Value = Vec<Vec<u8>>> {
15    prop::collection::vec(
16        prop::collection::vec(any::<u8>(), 0..=max_len),
17        0..=max_items,
18    )
19}
20
21/// A lane layout: 1..=`max_lanes` lanes with dense ids, each serving a
22/// distinct partition drawn from `0..max_partitions`. Suitable for
23/// [`SourceHandle::assign_lanes`](crate::SourceHandle::assign_lanes).
24pub fn lane_layout(
25    max_lanes: usize,
26    max_partitions: u32,
27) -> impl Strategy<Value = Vec<(LaneId, PartitionId)>> {
28    let all: Vec<u32> = (0..max_partitions).collect();
29    prop::sample::subsequence(all, 1..=max_lanes.min(max_partitions as usize)).prop_map(
30        |partitions| {
31            partitions
32                .into_iter()
33                .enumerate()
34                .map(|(i, p)| {
35                    (
36                        LaneId(u32::try_from(i).expect("dense lane id")),
37                        PartitionId(p),
38                    )
39                })
40                .collect()
41        },
42    )
43}
44
45/// One step of a generated source script. Indices are relative to a
46/// [`lane_layout`]; the consuming test resolves them and skips steps that
47/// are invalid in its current state (e.g. revoking an already-revoked
48/// lane).
49#[derive(Clone, Debug)]
50pub enum ScriptOp {
51    /// Push a payload to the partition at this layout index.
52    Push {
53        /// Index into the layout.
54        partition_index: usize,
55        /// Payload bytes.
56        payload: Vec<u8>,
57    },
58    /// Revoke the lane at this layout index.
59    Revoke {
60        /// Index into the layout.
61        lane_index: usize,
62    },
63    /// Re-assign the lane at this layout index (after a revoke).
64    Assign {
65        /// Index into the layout.
66        lane_index: usize,
67    },
68}
69
70/// A push-heavy script of up to `max_ops` steps over `layout_len` lanes,
71/// with payloads up to `max_payload_len` bytes.
72pub fn source_script(
73    max_ops: usize,
74    layout_len: usize,
75    max_payload_len: usize,
76) -> impl Strategy<Value = Vec<ScriptOp>> {
77    let op = prop_oneof![
78        8 => (0..layout_len, prop::collection::vec(any::<u8>(), 0..=max_payload_len))
79            .prop_map(|(partition_index, payload)| ScriptOp::Push { partition_index, payload }),
80        1 => (0..layout_len).prop_map(|lane_index| ScriptOp::Revoke { lane_index }),
81        1 => (0..layout_len).prop_map(|lane_index| ScriptOp::Assign { lane_index }),
82    ];
83    prop::collection::vec(op, 0..=max_ops)
84}
85
86/// Up to `max` scripted write outcomes, weighted towards success:
87/// 8 ok : 2 retryable : 1 delayed-ok (1..50 ms) : 1 fatal.
88pub fn write_outcomes(max: usize) -> impl Strategy<Value = Vec<WriteOutcome>> {
89    let outcome = prop_oneof![
90        8 => Just(WriteOutcome::ok()),
91        2 => Just(WriteOutcome::retryable("scripted retryable failure")),
92        1 => (1u64..50)
93            .prop_map(|ms| WriteOutcome::ok().after(Duration::from_millis(ms))),
94        1 => Just(WriteOutcome::fatal("scripted fatal failure")),
95    ];
96    prop::collection::vec(outcome, 0..=max)
97}