1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
#![allow(clippy::pedantic)]
use tenshift_core::error::Result;
use tenshift_core::pipeline::ErrorPolicy;
use tenshift_core::sample::{Sample, Tensor};
use tenshift_core::source::{Source, SourceIterator};
use tenshift_core::Pipeline;
/// A maliciously corrupted source simulating a disk error where a file chunk is surprisingly short.
/// We test if our pipeline gracefully aborts Collation instead of panic/segfault.
struct TruncatedFileSource {
total: usize,
}
impl Source for TruncatedFileSource {
fn open(&self) -> Result<Box<dyn SourceIterator>> {
Ok(Box::new(TruncatedIterator {
total: self.total,
idx: 0,
}))
}
fn name(&self) -> &str {
"trunc"
}
}
struct TruncatedIterator {
total: usize,
idx: usize,
}
impl SourceIterator for TruncatedIterator {
fn next_sample(&mut self) -> Option<Result<Sample>> {
if self.idx >= self.total {
return None;
}
let mut s = Sample::new();
// The first 4 batches are pristine.
// The 5th item has an unexpected sequence length (short read)
if self.idx == 4 {
// Emits 1 byte instead of 100 bytes!
s.insert("data", Tensor::i64(&[1], vec![1]));
} else {
// Normal 100-byte buffer
s.insert("data", Tensor::i64(&[0; 100], vec![100]));
}
self.idx += 1;
Some(Ok(s))
}
}
#[test]
fn test_sqlite_short_read_collation_abort() {
// When Collate encounters the 5th item with the wrong size, it must NOT panic.
// It must return an explicit `CollateFailed` error, bubbling up logically without breaking bounds.
let pipeline = Pipeline::from_source(TruncatedFileSource { total: 10 })
.batch(2)
.prefetch(1)
.workers(1)
.on_error(ErrorPolicy::Skip); // Automatically skip failed batches
let mut success_count = 0;
// Batch 1 (idx 0,1) -> success
// Batch 2 (idx 2,3) -> success
// Batch 3 (idx 4,5) -> Collate failure because idx 4 is truncated. Batch 3 is SKIPPED.
// Batch 4 (idx 6,7) -> success
// Batch 5 (idx 8,9) -> success
for _batch in pipeline.start().expect("init") {
success_count += 1;
}
// Since batch 3 fails collation due to mismatched schema (shape[1] vs shape[100]),
// it perfectly skips it without crashing the process.
assert_eq!(
success_count, 4,
"Pipeline crashed instead of catching short read collation"
);
}