1use crate::sink::WriteOutcome;
8use proptest::prelude::*;
9use spate_core::record::PartitionId;
10use spate_core::source::LaneId;
11use std::time::Duration;
12
13pub 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
21pub 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#[derive(Clone, Debug)]
50pub enum ScriptOp {
51 Push {
53 partition_index: usize,
55 payload: Vec<u8>,
57 },
58 Revoke {
60 lane_index: usize,
62 },
63 Assign {
65 lane_index: usize,
67 },
68}
69
70pub 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
86pub 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}