1use std::collections::HashMap;
11use std::pin::Pin;
12use std::sync::{Arc, Mutex};
13
14use faucet_core::write_mode::WriteMode;
15use faucet_core::{FaucetError, Sink, Source, StreamPage, Value, async_trait};
16use futures_core::Stream;
17use serde_json::json;
18
19pub struct CountingSource {
30 total: usize,
31 batch: usize,
32 resumable: bool,
33 start: Arc<Mutex<usize>>,
34}
35
36impl CountingSource {
37 pub fn new(total: usize, batch: usize) -> Self {
39 Self {
40 total,
41 batch,
42 resumable: true,
43 start: Arc::new(Mutex::new(0)),
44 }
45 }
46
47 pub fn non_resumable(total: usize, batch: usize) -> Self {
50 Self {
51 total,
52 batch,
53 resumable: false,
54 start: Arc::new(Mutex::new(0)),
55 }
56 }
57}
58
59#[async_trait]
60impl Source for CountingSource {
61 async fn fetch_with_context(
62 &self,
63 _context: &HashMap<String, Value>,
64 ) -> Result<Vec<Value>, FaucetError> {
65 let start = *self.start.lock().unwrap();
66 Ok((start..self.total).map(|i| json!({ "n": i })).collect())
67 }
68
69 fn stream_pages<'a>(
70 &'a self,
71 _context: &'a HashMap<String, Value>,
72 _batch_size: usize,
73 ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
74 let batch = if self.batch == 0 {
79 self.total.max(1)
80 } else {
81 self.batch
82 };
83 let total = self.total;
84 let start = (*self.start.lock().unwrap()).min(total);
85 Box::pin(async_stream::try_stream! {
86 let mut n = start;
87 if n >= total {
88 yield StreamPage { records: Vec::new(), bookmark: Some(json!({ "n": total })) };
91 return;
92 }
93 while n < total {
94 let end = (n + batch).min(total);
95 let records: Vec<Value> = (n..end).map(|i| json!({ "n": i })).collect();
96 n = end;
97 let bookmark = if n >= total { Some(json!({ "n": total })) } else { None };
98 yield StreamPage { records, bookmark };
99 }
100 })
101 }
102
103 fn connector_name(&self) -> &'static str {
104 "counting-source"
105 }
106
107 fn state_key(&self) -> Option<String> {
108 Some("conformance:counting".to_string())
109 }
110
111 async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
112 if !self.resumable {
113 return Ok(());
114 }
115 if let Some(n) = bookmark.get("n").and_then(|v| v.as_u64()) {
116 *self.start.lock().unwrap() = n as usize;
117 }
118 Ok(())
119 }
120}
121
122pub struct FailingSource;
126
127#[async_trait]
128impl Source for FailingSource {
129 async fn fetch_with_context(
130 &self,
131 _context: &HashMap<String, Value>,
132 ) -> Result<Vec<Value>, FaucetError> {
133 Err(FaucetError::Source(
134 "unreachable endpoint (test double)".to_string(),
135 ))
136 }
137
138 fn connector_name(&self) -> &'static str {
139 "failing-source"
140 }
141}
142
143pub struct PanickingSource;
147
148#[async_trait]
149impl Source for PanickingSource {
150 async fn fetch_with_context(
151 &self,
152 _context: &HashMap<String, Value>,
153 ) -> Result<Vec<Value>, FaucetError> {
154 panic!("connector bug: unwrap() on a None value");
155 }
156
157 fn connector_name(&self) -> &'static str {
158 "panicking-source"
159 }
160}
161
162#[derive(Clone, Default)]
173pub struct TestSink {
174 key_field: Option<String>,
175 idempotent: bool,
176 keyed: Arc<Mutex<HashMap<String, Value>>>,
177 appended: Arc<Mutex<Vec<Value>>>,
178 tokens: Arc<Mutex<HashMap<String, String>>>,
179 write_calls: Arc<Mutex<usize>>,
180}
181
182impl TestSink {
183 pub fn new() -> Self {
185 Self::default()
186 }
187
188 pub fn keyed(key_field: impl Into<String>) -> Self {
190 Self {
191 key_field: Some(key_field.into()),
192 ..Self::default()
193 }
194 }
195
196 pub fn idempotent(key_field: impl Into<String>) -> Self {
199 Self {
200 key_field: Some(key_field.into()),
201 idempotent: true,
202 ..Self::default()
203 }
204 }
205
206 pub fn len(&self) -> usize {
208 if self.key_field.is_some() {
209 self.keyed.lock().unwrap().len()
210 } else {
211 self.appended.lock().unwrap().len()
212 }
213 }
214
215 pub fn is_empty(&self) -> bool {
217 self.len() == 0
218 }
219
220 pub fn total_written(&self) -> usize {
223 *self.write_calls.lock().unwrap()
224 }
225}
226
227#[async_trait]
228impl Sink for TestSink {
229 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
230 *self.write_calls.lock().unwrap() += records.len();
231 match &self.key_field {
232 Some(field) => {
233 let mut map = self.keyed.lock().unwrap();
234 for r in records {
235 let key = r.get(field).map(|v| v.to_string()).ok_or_else(|| {
236 FaucetError::Sink(format!("record missing key `{field}`"))
237 })?;
238 map.insert(key, r.clone());
239 }
240 }
241 None => self
242 .appended
243 .lock()
244 .unwrap()
245 .extend(records.iter().cloned()),
246 }
247 Ok(records.len())
248 }
249
250 fn supports_idempotent_writes(&self) -> bool {
251 self.idempotent
252 }
253
254 fn dedups_by_key(&self) -> bool {
255 self.key_field.is_some()
256 }
257
258 fn supported_write_modes(&self) -> &'static [WriteMode] {
259 if self.key_field.is_some() {
260 &[WriteMode::Append, WriteMode::Upsert, WriteMode::Delete]
261 } else {
262 &[WriteMode::Append]
263 }
264 }
265
266 async fn write_batch_idempotent(
267 &self,
268 records: &[Value],
269 scope: &str,
270 token: &str,
271 ) -> Result<usize, FaucetError> {
272 self.tokens
276 .lock()
277 .unwrap()
278 .insert(scope.to_string(), token.to_string());
279 self.write_batch(records).await
280 }
281
282 async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
283 Ok(self.tokens.lock().unwrap().get(scope).cloned())
284 }
285
286 fn connector_name(&self) -> &'static str {
287 "test-sink"
288 }
289}
290
291#[derive(Clone, Default)]
295pub struct LyingIdempotentSink {
296 appended: Arc<Mutex<Vec<Value>>>,
297}
298
299impl LyingIdempotentSink {
300 pub fn new() -> Self {
302 Self::default()
303 }
304 pub fn len(&self) -> usize {
306 self.appended.lock().unwrap().len()
307 }
308 pub fn is_empty(&self) -> bool {
310 self.len() == 0
311 }
312}
313
314#[async_trait]
315impl Sink for LyingIdempotentSink {
316 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
317 self.appended
318 .lock()
319 .unwrap()
320 .extend(records.iter().cloned());
321 Ok(records.len())
322 }
323
324 fn supports_idempotent_writes(&self) -> bool {
325 true }
327
328 fn connector_name(&self) -> &'static str {
329 "lying-idempotent-sink"
330 }
331}
332
333#[derive(Clone, Default)]
337pub struct LyingKeyedSink {
338 appended: Arc<Mutex<Vec<Value>>>,
339}
340
341impl LyingKeyedSink {
342 pub fn new() -> Self {
344 Self::default()
345 }
346 pub fn len(&self) -> usize {
348 self.appended.lock().unwrap().len()
349 }
350 pub fn is_empty(&self) -> bool {
352 self.len() == 0
353 }
354}
355
356#[async_trait]
357impl Sink for LyingKeyedSink {
358 async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
359 self.appended
360 .lock()
361 .unwrap()
362 .extend(records.iter().cloned());
363 Ok(records.len())
364 }
365
366 fn dedups_by_key(&self) -> bool {
367 true }
369
370 fn supported_write_modes(&self) -> &'static [WriteMode] {
371 &[WriteMode::Append, WriteMode::Upsert]
372 }
373
374 fn connector_name(&self) -> &'static str {
375 "lying-keyed-sink"
376 }
377}
378
379#[cfg(test)]
380mod tests {
381 use super::*;
382 use serde_json::json;
383 use std::collections::HashMap;
384
385 #[tokio::test]
386 async fn counting_source_resumes_and_ignores_when_non_resumable() {
387 let s = CountingSource::new(5, 2);
388 assert_eq!(s.state_key().as_deref(), Some("conformance:counting"));
389 assert_eq!(s.connector_name(), "counting-source");
390 assert_eq!(
391 s.fetch_with_context(&HashMap::new()).await.unwrap().len(),
392 5
393 );
394 s.apply_start_bookmark(json!({ "n": 5 })).await.unwrap();
396 assert!(
397 s.fetch_with_context(&HashMap::new())
398 .await
399 .unwrap()
400 .is_empty()
401 );
402
403 let nr = CountingSource::non_resumable(5, 2);
405 nr.apply_start_bookmark(json!({ "n": 5 })).await.unwrap();
406 assert_eq!(
407 nr.fetch_with_context(&HashMap::new()).await.unwrap().len(),
408 5
409 );
410 }
411
412 #[tokio::test]
413 async fn test_sink_accessors() {
414 let s = TestSink::new();
415 assert!(s.is_empty());
416 s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
417 assert!(!s.is_empty());
418 assert_eq!(s.len(), 1);
419 assert_eq!(s.total_written(), 1);
420 assert_eq!(s.connector_name(), "test-sink");
421 }
422
423 #[tokio::test]
424 async fn lying_idempotent_sink_never_persists_a_token() {
425 let s = LyingIdempotentSink::new();
426 assert!(s.is_empty());
427 assert!(s.supports_idempotent_writes());
428 assert_eq!(s.connector_name(), "lying-idempotent-sink");
429 s.write_batch_idempotent(&[json!({ "id": 1 })], "scope", "00000000000000000001")
430 .await
431 .unwrap();
432 assert_eq!(s.len(), 1);
433 assert!(s.last_committed_token("scope").await.unwrap().is_none());
434 }
435
436 #[tokio::test]
437 async fn lying_keyed_sink_appends_duplicates() {
438 let s = LyingKeyedSink::new();
439 assert!(s.is_empty());
440 assert!(s.dedups_by_key());
441 assert!(s.supported_write_modes().contains(&WriteMode::Upsert));
442 assert_eq!(s.connector_name(), "lying-keyed-sink");
443 s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
444 s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
445 assert_eq!(s.len(), 2, "lying keyed sink does not dedup");
446 }
447
448 #[tokio::test]
449 async fn failing_and_panicking_source_labels() {
450 assert_eq!(FailingSource.connector_name(), "failing-source");
451 assert_eq!(PanickingSource.connector_name(), "panicking-source");
452 assert!(FailingSource.fetch_all().await.is_err());
453 }
454}