1use spate_core::checkpoint::AckRef;
4use spate_core::deser::{Deserializer, EmitRecord, Owned};
5use spate_core::error::DeserError;
6use spate_core::record::{Flow, RawPayload, Record};
7
8#[derive(Clone, Debug)]
9enum Mode {
10 Passthrough,
11 SplitOn(u8),
12 FailOnPrefix(Vec<u8>),
13}
14
15#[derive(Clone, Debug)]
29pub struct TestDeserializer {
30 mode: Mode,
31}
32
33impl TestDeserializer {
34 #[must_use]
36 pub fn passthrough() -> Self {
37 TestDeserializer {
38 mode: Mode::Passthrough,
39 }
40 }
41
42 #[must_use]
44 pub fn split_on(delimiter: u8) -> Self {
45 TestDeserializer {
46 mode: Mode::SplitOn(delimiter),
47 }
48 }
49
50 #[must_use]
52 pub fn fail_on_prefix(marker: impl Into<Vec<u8>>) -> Self {
53 TestDeserializer {
54 mode: Mode::FailOnPrefix(marker.into()),
55 }
56 }
57}
58
59impl Deserializer<Owned<Vec<u8>>> for TestDeserializer {
60 fn deserialize<'buf>(
61 &mut self,
62 raw: &RawPayload<'buf>,
63 ack: &AckRef,
64 out: &mut dyn EmitRecord<'buf, Vec<u8>>,
65 ) -> Result<(), DeserError> {
66 let meta = raw.meta();
67 let mut emit = |payload: Vec<u8>| {
68 out.emit(Record {
69 payload,
70 meta,
71 ack: ack.clone(),
72 })
73 };
74 match &self.mode {
75 Mode::Passthrough => {
76 let _ = emit(raw.bytes.to_vec());
77 }
78 Mode::SplitOn(delim) => {
79 for segment in raw.bytes.split(|b| b == delim) {
80 if emit(segment.to_vec()) == Flow::Blocked {
81 break;
82 }
83 }
84 }
85 Mode::FailOnPrefix(marker) => {
86 if raw.bytes.starts_with(marker) {
87 return Err(DeserError::Malformed {
88 reason: format!("payload starts with test marker {marker:?}"),
89 });
90 }
91 let _ = emit(raw.bytes.to_vec());
92 }
93 }
94 Ok(())
95 }
96}
97
98#[derive(Debug, Default)]
103pub struct EmitCollector<T> {
104 pub records: Vec<Record<T>>,
106 pub rejected: usize,
108 block_after: Option<usize>,
109}
110
111impl<T> EmitCollector<T> {
112 #[must_use]
114 pub fn new() -> Self {
115 EmitCollector {
116 records: Vec::new(),
117 rejected: 0,
118 block_after: None,
119 }
120 }
121
122 #[must_use]
125 pub fn blocking_after(n: usize) -> Self {
126 EmitCollector {
127 records: Vec::new(),
128 rejected: 0,
129 block_after: Some(n),
130 }
131 }
132
133 #[must_use]
135 pub fn payloads(&self) -> Vec<T>
136 where
137 T: Clone,
138 {
139 self.records.iter().map(|r| r.payload.clone()).collect()
140 }
141}
142
143impl<'buf, T> EmitRecord<'buf, T> for EmitCollector<T> {
144 fn emit(&mut self, rec: Record<T>) -> Flow {
145 if let Some(cap) = self.block_after
146 && self.records.len() >= cap
147 {
148 self.rejected += 1;
149 return Flow::Blocked;
150 }
151 self.records.push(rec);
152 Flow::Continue
153 }
154}