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::Initialized { .. } => {
138                return Err(ConnectorError::ConfigurationError(
139                    "mock source has no sealed topology startup contract".into(),
140                ));
141            }
142            crate::connector::SourcePosition::Resume { checkpoint, .. } => checkpoint
143                .get_offset("records")
144                .ok_or_else(|| {
145                    ConnectorError::ConfigurationError(
146                        "mock source checkpoint is missing 'records'".into(),
147                    )
148                })?
149                .parse::<u64>()
150                .map_err(|error| {
151                    ConnectorError::ConfigurationError(format!(
152                        "invalid mock source record cursor: {error}"
153                    ))
154                })?,
155        };
156        let consumed_batches = if self.batch_size == 0 {
157            0
158        } else {
159            records / self.batch_size as u64
160        };
161        self.records_produced.store(records, Ordering::Relaxed);
162        self.batches_remaining.store(
163            self.initial_batches.saturating_sub(consumed_batches),
164            Ordering::Relaxed,
165        );
166        self.is_open
167            .store(true, std::sync::atomic::Ordering::Relaxed);
168        Ok(())
169    }
170
171    async fn poll_batch(
172        &mut self,
173        _max_records: usize,
174    ) -> Result<Option<SourceBatch>, ConnectorError> {
175        let remaining = self.batches_remaining.load(Ordering::Relaxed);
176        if remaining == 0 {
177            return Ok(None);
178        }
179        self.batches_remaining.fetch_sub(1, Ordering::Relaxed);
180
181        let batch = mock_batch(self.batch_size);
182        self.records_produced
183            .fetch_add(self.batch_size as u64, Ordering::Relaxed);
184        Ok(Some(SourceBatch::new(batch)))
185    }
186
187    fn schema(&self) -> SchemaRef {
188        self.schema.clone()
189    }
190
191    fn checkpoint(&self) -> SourceCheckpoint {
192        let mut cp = SourceCheckpoint::new();
193        cp.set_offset(
194            "records",
195            self.records_produced.load(Ordering::Relaxed).to_string(),
196        );
197        cp
198    }
199
200    async fn close(&mut self) -> Result<(), ConnectorError> {
201        self.is_open
202            .store(false, std::sync::atomic::Ordering::Relaxed);
203        Ok(())
204    }
205
206    async fn notify_epoch_committed(
207        &mut self,
208        epoch: u64,
209        _checkpoint: &SourceCheckpoint,
210    ) -> Result<(), ConnectorError> {
211        self.committed_epochs.lock().push(epoch);
212        Ok(())
213    }
214}
215
216/// Mock sink connector for testing.
217///
218/// Stores all written batches in memory for inspection.
219#[derive(Debug)]
220pub struct MockSinkConnector {
221    schema: SchemaRef,
222    written: Arc<Mutex<Vec<RecordBatch>>>,
223    records_written: AtomicU64,
224    is_open: std::sync::atomic::AtomicBool,
225}
226
227impl MockSinkConnector {
228    /// Creates a new mock sink.
229    #[must_use]
230    pub fn new() -> Self {
231        Self {
232            schema: mock_schema(),
233            written: Arc::new(Mutex::new(Vec::new())),
234            records_written: AtomicU64::new(0),
235            is_open: std::sync::atomic::AtomicBool::new(false),
236        }
237    }
238
239    /// Returns the number of batches written.
240    #[must_use]
241    pub fn batch_count(&self) -> usize {
242        self.written.lock().len()
243    }
244
245    /// Returns the total number of records written.
246    #[must_use]
247    pub fn records_written(&self) -> u64 {
248        self.records_written.load(Ordering::Relaxed)
249    }
250
251    /// Returns a clone of all written batches.
252    #[must_use]
253    pub fn written_batches(&self) -> Vec<RecordBatch> {
254        self.written.lock().clone()
255    }
256}
257
258impl Default for MockSinkConnector {
259    fn default() -> Self {
260        Self::new()
261    }
262}
263
264#[async_trait]
265impl SinkConnector for MockSinkConnector {
266    fn contract(&self, _config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
267        Ok(SinkContract::new(
268            SinkConsistency::Ephemeral,
269            SinkTopology::NodeLocalEgress,
270            SinkInputMode::AppendOnly,
271        ))
272    }
273
274    async fn open(&mut self, _config: &ConnectorConfig) -> Result<(), ConnectorError> {
275        self.is_open
276            .store(true, std::sync::atomic::Ordering::Relaxed);
277        Ok(())
278    }
279
280    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
281        let num_rows = batch.num_rows();
282        let bytes = batch.get_array_memory_size() as u64;
283        self.written.lock().push(batch.clone());
284        self.records_written
285            .fetch_add(num_rows as u64, Ordering::Relaxed);
286        Ok(WriteResult::new(num_rows, bytes))
287    }
288
289    fn schema(&self) -> SchemaRef {
290        self.schema.clone()
291    }
292
293    fn suggested_write_timeout(&self) -> Duration {
294        Duration::from_secs(60)
295    }
296
297    async fn close(&mut self) -> Result<(), ConnectorError> {
298        self.is_open
299            .store(false, std::sync::atomic::Ordering::Relaxed);
300        Ok(())
301    }
302}
303
304/// Registers a mock source connector with the registry.
305///
306/// # Errors
307/// Returns an error when the mock source name is already registered.
308pub fn register_mock_source(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
309    registry.register_source(
310        "mock",
311        ConnectorInfo {
312            schema_capabilities: crate::schema::resolution::SchemaCapabilities::declared(false),
313            name: "mock".to_string(),
314            display_name: "Mock Source".to_string(),
315            version: "0.1.0".to_string(),
316            is_source: true,
317            is_sink: false,
318            config_keys: vec![],
319        },
320        Arc::new(|_: Option<&Arc<prometheus::Registry>>| Ok(Box::new(MockSourceConnector::new()))),
321    )
322}
323
324/// Registers a mock sink connector with the registry.
325///
326/// # Errors
327/// Returns an error when the mock sink name is already registered.
328pub fn register_mock_sink(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
329    registry.register_sink(
330        "mock",
331        ConnectorInfo {
332            schema_capabilities: crate::schema::resolution::SchemaCapabilities::declared(true),
333            name: "mock".to_string(),
334            display_name: "Mock Sink".to_string(),
335            version: "0.1.0".to_string(),
336            is_source: false,
337            is_sink: true,
338            config_keys: vec![],
339        },
340        Arc::new(|_config, _registry| Ok(Box::new(MockSinkConnector::new()))),
341    )
342}
343
344#[cfg(test)]
345mod tests;
laminar_connectors/testing/mod.rs

laminar_connectors/testing/
mod.rs
