laminar_connectors/testing/
mod.rs

1//! Testing utilities for connector implementations.
2//!
3//! Provides mock connectors and helper functions for testing
4//! the connector SDK and concrete connector implementations.
5
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::Arc;
8use std::time::Duration;
9
10use arrow_array::{Int64Array, RecordBatch, StringArray};
11use arrow_schema::{DataType, Field, Schema, SchemaRef};
12use async_trait::async_trait;
13use parking_lot::Mutex;
14
15use crate::checkpoint::SourceCheckpoint;
16use crate::config::{ConnectorConfig, ConnectorInfo};
17use crate::connector::{
18    SinkConnector, SinkConsistency, SinkContract, SinkInputMode, SinkTopology, SourceBatch,
19    SourceConnector, SourceConsistency, SourceContract, SourceInputMode, SourceTopology,
20    WriteResult,
21};
22use crate::error::ConnectorError;
23use crate::registry::ConnectorRegistry;
24
25/// Creates a test schema with `id` (Int64) and `value` (Utf8) columns.
26#[must_use]
27pub fn mock_schema() -> SchemaRef {
28    Arc::new(Schema::new(vec![
29        Field::new("id", DataType::Int64, false),
30        Field::new("value", DataType::Utf8, false),
31    ]))
32}
33
34/// Creates a test `RecordBatch` with `n` rows.
35///
36/// # Panics
37///
38/// Panics if the batch cannot be created (should not happen with valid inputs).
39#[must_use]
40pub fn mock_batch(n: usize) -> RecordBatch {
41    #[allow(clippy::cast_possible_wrap)]
42    let ids: Vec<i64> = (0..n as i64).collect();
43    let values: Vec<String> = (0..n).map(|i| format!("value_{i}")).collect();
44    let value_refs: Vec<&str> = values.iter().map(String::as_str).collect();
45
46    RecordBatch::try_new(
47        mock_schema(),
48        vec![
49            Arc::new(Int64Array::from(ids)),
50            Arc::new(StringArray::from(value_refs)),
51        ],
52    )
53    .unwrap()
54}
55
56/// Mock source connector for testing.
57///
58/// Returns a configurable number of batches, then returns `None`.
59#[derive(Debug)]
60pub struct MockSourceConnector {
61    schema: SchemaRef,
62    initial_batches: u64,
63    batches_remaining: AtomicU64,
64    batch_size: usize,
65    records_produced: AtomicU64,
66    is_open: std::sync::atomic::AtomicBool,
67    committed_epochs: Arc<Mutex<Vec<u64>>>,
68}
69
70impl MockSourceConnector {
71    /// Creates a new mock source that produces 10 batches of 5 records.
72    #[must_use]
73    pub fn new() -> Self {
74        Self {
75            schema: mock_schema(),
76            initial_batches: 10,
77            batches_remaining: AtomicU64::new(10),
78            batch_size: 5,
79            records_produced: AtomicU64::new(0),
80            is_open: std::sync::atomic::AtomicBool::new(false),
81            committed_epochs: Arc::new(Mutex::new(Vec::new())),
82        }
83    }
84
85    /// Creates a mock source with custom batch count and size.
86    #[must_use]
87    pub fn with_batches(count: u64, batch_size: usize) -> Self {
88        Self {
89            schema: mock_schema(),
90            initial_batches: count,
91            batches_remaining: AtomicU64::new(count),
92            batch_size,
93            records_produced: AtomicU64::new(0),
94            is_open: std::sync::atomic::AtomicBool::new(false),
95            committed_epochs: Arc::new(Mutex::new(Vec::new())),
96        }
97    }
98
99    /// Returns the total records produced.
100    #[must_use]
101    pub fn records_produced(&self) -> u64 {
102        self.records_produced.load(Ordering::Relaxed)
103    }
104
105    /// Returns a handle for inspecting epochs reported to
106    /// `notify_epoch_committed` from outside the connector (the
107    /// connector itself is moved into a background task when run via
108    /// the pipeline).
109    #[must_use]
110    pub fn committed_epochs_handle(&self) -> Arc<Mutex<Vec<u64>>> {
111        Arc::clone(&self.committed_epochs)
112    }
113}
114
115impl Default for MockSourceConnector {
116    fn default() -> Self {
117        Self::new()
118    }
119}
120
121#[async_trait]
122impl SourceConnector for MockSourceConnector {
123    fn contract(&self, _config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
124        Ok(SourceContract::new(
125            SourceConsistency::Replayable,
126            SourceTopology::Singleton,
127            SourceInputMode::AppendOnly,
128        ))
129    }
130
131    async fn start(
132        &mut self,
133        request: crate::connector::SourceStart,
134    ) -> Result<(), ConnectorError> {
135        let records = match request.into_parts().1 {
136            crate::connector::SourcePosition::Initial => 0,
137            crate::connector::SourcePosition::Resume { checkpoint, .. } => checkpoint
138                .get_offset("records")
139                .ok_or_else(|| {
140                    ConnectorError::ConfigurationError(
141                        "mock source checkpoint is missing 'records'".into(),
142                    )
143                })?
144                .parse::<u64>()
145                .map_err(|error| {
146                    ConnectorError::ConfigurationError(format!(
147                        "invalid mock source record cursor: {error}"
148                    ))
149                })?,
150        };
151        let consumed_batches = if self.batch_size == 0 {
152            0
153        } else {
154            records / self.batch_size as u64
155        };
156        self.records_produced.store(records, Ordering::Relaxed);
157        self.batches_remaining.store(
158            self.initial_batches.saturating_sub(consumed_batches),
159            Ordering::Relaxed,
160        );
161        self.is_open
162            .store(true, std::sync::atomic::Ordering::Relaxed);
163        Ok(())
164    }
165
166    async fn poll_batch(
167        &mut self,
168        _max_records: usize,
169    ) -> Result<Option<SourceBatch>, ConnectorError> {
170        let remaining = self.batches_remaining.load(Ordering::Relaxed);
171        if remaining == 0 {
172            return Ok(None);
173        }
174        self.batches_remaining.fetch_sub(1, Ordering::Relaxed);
175
176        let batch = mock_batch(self.batch_size);
177        self.records_produced
178            .fetch_add(self.batch_size as u64, Ordering::Relaxed);
179        Ok(Some(SourceBatch::new(batch)))
180    }
181
182    fn schema(&self) -> SchemaRef {
183        self.schema.clone()
184    }
185
186    fn checkpoint(&self) -> SourceCheckpoint {
187        let mut cp = SourceCheckpoint::new();
188        cp.set_offset(
189            "records",
190            self.records_produced.load(Ordering::Relaxed).to_string(),
191        );
192        cp
193    }
194
195    async fn close(&mut self) -> Result<(), ConnectorError> {
196        self.is_open
197            .store(false, std::sync::atomic::Ordering::Relaxed);
198        Ok(())
199    }
200
201    async fn notify_epoch_committed(
202        &mut self,
203        epoch: u64,
204        _checkpoint: &SourceCheckpoint,
205    ) -> Result<(), ConnectorError> {
206        self.committed_epochs.lock().push(epoch);
207        Ok(())
208    }
209}
210
211/// Mock sink connector for testing.
212///
213/// Stores all written batches in memory for inspection.
214#[derive(Debug)]
215pub struct MockSinkConnector {
216    schema: SchemaRef,
217    written: Arc<Mutex<Vec<RecordBatch>>>,
218    records_written: AtomicU64,
219    is_open: std::sync::atomic::AtomicBool,
220}
221
222impl MockSinkConnector {
223    /// Creates a new mock sink.
224    #[must_use]
225    pub fn new() -> Self {
226        Self {
227            schema: mock_schema(),
228            written: Arc::new(Mutex::new(Vec::new())),
229            records_written: AtomicU64::new(0),
230            is_open: std::sync::atomic::AtomicBool::new(false),
231        }
232    }
233
234    /// Returns the number of batches written.
235    #[must_use]
236    pub fn batch_count(&self) -> usize {
237        self.written.lock().len()
238    }
239
240    /// Returns the total number of records written.
241    #[must_use]
242    pub fn records_written(&self) -> u64 {
243        self.records_written.load(Ordering::Relaxed)
244    }
245
246    /// Returns a clone of all written batches.
247    #[must_use]
248    pub fn written_batches(&self) -> Vec<RecordBatch> {
249        self.written.lock().clone()
250    }
251}
252
253impl Default for MockSinkConnector {
254    fn default() -> Self {
255        Self::new()
256    }
257}
258
259#[async_trait]
260impl SinkConnector for MockSinkConnector {
261    fn contract(&self, _config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
262        Ok(SinkContract::new(
263            SinkConsistency::Ephemeral,
264            SinkTopology::NodeLocalEgress,
265            SinkInputMode::AppendOnly,
266        ))
267    }
268
269    async fn open(&mut self, _config: &ConnectorConfig) -> Result<(), ConnectorError> {
270        self.is_open
271            .store(true, std::sync::atomic::Ordering::Relaxed);
272        Ok(())
273    }
274
275    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
276        let num_rows = batch.num_rows();
277        let bytes = batch.get_array_memory_size() as u64;
278        self.written.lock().push(batch.clone());
279        self.records_written
280            .fetch_add(num_rows as u64, Ordering::Relaxed);
281        Ok(WriteResult::new(num_rows, bytes))
282    }
283
284    fn schema(&self) -> SchemaRef {
285        self.schema.clone()
286    }
287
288    fn suggested_write_timeout(&self) -> Duration {
289        Duration::from_secs(60)
290    }
291
292    async fn close(&mut self) -> Result<(), ConnectorError> {
293        self.is_open
294            .store(false, std::sync::atomic::Ordering::Relaxed);
295        Ok(())
296    }
297}
298
299/// Registers a mock source connector with the registry.
300///
301/// # Errors
302/// Returns an error when the mock source name is already registered.
303pub fn register_mock_source(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
304    registry.register_source(
305        "mock",
306        ConnectorInfo {
307            name: "mock".to_string(),
308            display_name: "Mock Source".to_string(),
309            version: "0.1.0".to_string(),
310            is_source: true,
311            is_sink: false,
312            config_keys: vec![],
313        },
314        Arc::new(|_: Option<&Arc<prometheus::Registry>>| Ok(Box::new(MockSourceConnector::new()))),
315    )
316}
317
318/// Registers a mock sink connector with the registry.
319///
320/// # Errors
321/// Returns an error when the mock sink name is already registered.
322pub fn register_mock_sink(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
323    registry.register_sink(
324        "mock",
325        ConnectorInfo {
326            name: "mock".to_string(),
327            display_name: "Mock Sink".to_string(),
328            version: "0.1.0".to_string(),
329            is_source: false,
330            is_sink: true,
331            config_keys: vec![],
332        },
333        Arc::new(|_config, _registry| Ok(Box::new(MockSinkConnector::new()))),
334    )
335}
336
337#[cfg(test)]
338mod tests;
laminar_connectors/testing/mod.rs

laminar_connectors/testing/
mod.rs
