use std::collections::HashMap;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use faucet_core::check::{CheckContext, CheckReport};
use faucet_core::drift::SchemaEvolution;
use faucet_core::write_mode::{DeleteMarker, WriteMode};
use faucet_core::{FaucetError, Sink, Source, StreamPage, Value, async_trait};
use futures_core::Stream;
use serde_json::json;
pub const DELETE_MARKER_FIELD: &str = "__op";
pub const DELETE_MARKER_VALUE: &str = "d";
pub struct CountingSource {
total: usize,
batch: usize,
resumable: bool,
start: Arc<Mutex<usize>>,
}
impl CountingSource {
pub fn new(total: usize, batch: usize) -> Self {
Self {
total,
batch,
resumable: true,
start: Arc::new(Mutex::new(0)),
}
}
pub fn non_resumable(total: usize, batch: usize) -> Self {
Self {
total,
batch,
resumable: false,
start: Arc::new(Mutex::new(0)),
}
}
}
#[async_trait]
impl Source for CountingSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
let start = *self.start.lock().unwrap();
Ok((start..self.total).map(|i| json!({ "n": i })).collect())
}
fn stream_pages<'a>(
&'a self,
_context: &'a HashMap<String, Value>,
_batch_size: usize,
) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
let batch = if self.batch == 0 {
self.total.max(1)
} else {
self.batch
};
let total = self.total;
let start = (*self.start.lock().unwrap()).min(total);
Box::pin(async_stream::try_stream! {
let mut n = start;
if n >= total {
yield StreamPage { records: Vec::new(), bookmark: Some(json!({ "n": total })) };
return;
}
while n < total {
let end = (n + batch).min(total);
let records: Vec<Value> = (n..end).map(|i| json!({ "n": i })).collect();
n = end;
let bookmark = if n >= total { Some(json!({ "n": total })) } else { None };
yield StreamPage { records, bookmark };
}
})
}
fn connector_name(&self) -> &'static str {
"counting-source"
}
fn state_key(&self) -> Option<String> {
Some("conformance:counting".to_string())
}
async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
if !self.resumable {
return Ok(());
}
if let Some(n) = bookmark.get("n").and_then(|v| v.as_u64()) {
*self.start.lock().unwrap() = n as usize;
}
Ok(())
}
}
pub struct FailingSource;
#[async_trait]
impl Source for FailingSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
Err(FaucetError::Source(
"unreachable endpoint (test double)".to_string(),
))
}
fn connector_name(&self) -> &'static str {
"failing-source"
}
}
pub struct PanickingSource;
#[async_trait]
impl Source for PanickingSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
panic!("connector bug: unwrap() on a None value");
}
fn connector_name(&self) -> &'static str {
"panicking-source"
}
}
#[derive(Clone, Default)]
pub struct TestSink {
key_field: Option<String>,
idempotent: bool,
delete_marker: Option<DeleteMarker>,
keyed: Arc<Mutex<HashMap<String, Value>>>,
appended: Arc<Mutex<Vec<Value>>>,
tokens: Arc<Mutex<HashMap<String, String>>>,
write_calls: Arc<Mutex<usize>>,
}
impl TestSink {
pub fn new() -> Self {
Self::default()
}
pub fn keyed(key_field: impl Into<String>) -> Self {
Self {
key_field: Some(key_field.into()),
..Self::default()
}
}
pub fn keyed_upsert(key_field: impl Into<String>) -> Self {
Self {
key_field: Some(key_field.into()),
delete_marker: Some(DeleteMarker {
field: DELETE_MARKER_FIELD.to_string(),
values: vec![DELETE_MARKER_VALUE.to_string()],
}),
..Self::default()
}
}
pub fn idempotent(key_field: impl Into<String>) -> Self {
Self {
key_field: Some(key_field.into()),
idempotent: true,
..Self::default()
}
}
pub fn len(&self) -> usize {
if self.key_field.is_some() {
self.keyed.lock().unwrap().len()
} else {
self.appended.lock().unwrap().len()
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn total_written(&self) -> usize {
*self.write_calls.lock().unwrap()
}
fn is_delete_marked(&self, record: &Value) -> bool {
match &self.delete_marker {
Some(dm) => record
.get(&dm.field)
.and_then(|v| v.as_str())
.is_some_and(|s| dm.values.iter().any(|m| m == s)),
None => false,
}
}
}
#[async_trait]
impl Sink for TestSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
*self.write_calls.lock().unwrap() += records.len();
match &self.key_field {
Some(field) => {
let mut map = self.keyed.lock().unwrap();
for r in records {
let key = r.get(field).map(|v| v.to_string()).ok_or_else(|| {
FaucetError::Sink(format!("record missing key `{field}`"))
})?;
if self.is_delete_marked(r) {
map.remove(&key);
} else {
map.insert(key, r.clone());
}
}
}
None => self
.appended
.lock()
.unwrap()
.extend(records.iter().cloned()),
}
Ok(records.len())
}
fn supports_idempotent_writes(&self) -> bool {
self.idempotent
}
fn dedups_by_key(&self) -> bool {
self.key_field.is_some()
}
fn supported_write_modes(&self) -> &'static [WriteMode] {
if self.key_field.is_some() {
&[WriteMode::Append, WriteMode::Upsert, WriteMode::Delete]
} else {
&[WriteMode::Append]
}
}
async fn write_batch_idempotent(
&self,
records: &[Value],
scope: &str,
token: &str,
) -> Result<usize, FaucetError> {
self.tokens
.lock()
.unwrap()
.insert(scope.to_string(), token.to_string());
self.write_batch(records).await
}
async fn last_committed_token(&self, scope: &str) -> Result<Option<String>, FaucetError> {
Ok(self.tokens.lock().unwrap().get(scope).cloned())
}
fn connector_name(&self) -> &'static str {
"test-sink"
}
}
#[derive(Clone, Default)]
pub struct LyingIdempotentSink {
appended: Arc<Mutex<Vec<Value>>>,
}
impl LyingIdempotentSink {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.appended.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl Sink for LyingIdempotentSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.appended
.lock()
.unwrap()
.extend(records.iter().cloned());
Ok(records.len())
}
fn supports_idempotent_writes(&self) -> bool {
true }
fn connector_name(&self) -> &'static str {
"lying-idempotent-sink"
}
}
#[derive(Clone, Default)]
pub struct LyingKeyedSink {
appended: Arc<Mutex<Vec<Value>>>,
}
impl LyingKeyedSink {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.appended.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl Sink for LyingKeyedSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
self.appended
.lock()
.unwrap()
.extend(records.iter().cloned());
Ok(records.len())
}
fn dedups_by_key(&self) -> bool {
true }
fn supported_write_modes(&self) -> &'static [WriteMode] {
&[WriteMode::Append, WriteMode::Upsert]
}
fn connector_name(&self) -> &'static str {
"lying-keyed-sink"
}
}
#[derive(Clone)]
pub struct EvolvingSink {
columns: Arc<Mutex<HashMap<String, Value>>>,
}
impl Default for EvolvingSink {
fn default() -> Self {
let mut cols = HashMap::new();
cols.insert("id".to_string(), json!({ "type": "integer" }));
Self {
columns: Arc::new(Mutex::new(cols)),
}
}
}
impl EvolvingSink {
pub fn new() -> Self {
Self::default()
}
pub fn column_count(&self) -> usize {
self.columns.lock().unwrap().len()
}
}
fn schema_from_columns(cols: &HashMap<String, Value>) -> Value {
let props: serde_json::Map<String, Value> =
cols.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
json!({ "type": "object", "properties": props })
}
#[async_trait]
impl Sink for EvolvingSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
Ok(records.len())
}
async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
Ok(Some(schema_from_columns(&self.columns.lock().unwrap())))
}
fn supports_schema_evolution(&self) -> bool {
true
}
async fn evolve_schema(&self, evolution: &SchemaEvolution) -> Result<(), FaucetError> {
let mut cols = self.columns.lock().unwrap();
for change in evolution.additions.iter().chain(&evolution.widenings) {
cols.insert(change.name.clone(), change.to.clone());
}
Ok(())
}
fn connector_name(&self) -> &'static str {
"evolving-sink"
}
}
#[derive(Clone, Default)]
pub struct NoOpEvolvingSink;
#[async_trait]
impl Sink for NoOpEvolvingSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
Ok(records.len())
}
async fn current_schema(&self) -> Result<Option<Value>, FaucetError> {
Ok(Some(json!({
"type": "object",
"properties": { "id": { "type": "integer" } }
})))
}
fn supports_schema_evolution(&self) -> bool {
true }
async fn evolve_schema(&self, _evolution: &SchemaEvolution) -> Result<(), FaucetError> {
Ok(()) }
fn connector_name(&self) -> &'static str {
"noop-evolving-sink"
}
}
pub struct MultiPageZeroSource {
total: usize,
page: usize,
}
impl MultiPageZeroSource {
pub fn new(total: usize) -> Self {
Self { total, page: 2 }
}
}
#[async_trait]
impl Source for MultiPageZeroSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
Ok((0..self.total).map(|i| json!({ "n": i })).collect())
}
fn stream_pages<'a>(
&'a self,
_context: &'a HashMap<String, Value>,
_batch_size: usize,
) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
let total = self.total;
let page = self.page.max(1);
Box::pin(async_stream::try_stream! {
let mut n = 0;
while n < total {
let end = (n + page).min(total);
let records: Vec<Value> = (n..end).map(|i| json!({ "n": i })).collect();
n = end;
let bookmark = if n >= total { Some(json!({ "n": total })) } else { None };
yield StreamPage { records, bookmark };
}
})
}
fn connector_name(&self) -> &'static str {
"multi-page-zero-source"
}
}
pub struct EmptyNameSource;
#[async_trait]
impl Source for EmptyNameSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
Ok(Vec::new())
}
fn connector_name(&self) -> &'static str {
"" }
}
pub struct ErringCheckSource;
#[async_trait]
impl Source for ErringCheckSource {
async fn fetch_with_context(
&self,
_context: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
Ok(Vec::new())
}
async fn check(&self, _ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
Err(FaucetError::Source(
"probe failed — but returned as Err instead of a Fail probe".to_string(),
))
}
fn connector_name(&self) -> &'static str {
"erring-check-source"
}
}
#[derive(Clone, Default)]
pub struct ErringCheckSink;
#[async_trait]
impl Sink for ErringCheckSink {
async fn write_batch(&self, records: &[Value]) -> Result<usize, FaucetError> {
Ok(records.len())
}
async fn check(&self, _ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
Err(FaucetError::Sink(
"probe failed — but returned as Err instead of a Fail probe".to_string(),
))
}
fn connector_name(&self) -> &'static str {
"erring-check-sink"
}
}
#[cfg(test)]
mod tests {
use super::*;
use faucet_core::drift::ColumnChange;
use futures::StreamExt;
use serde_json::json;
use std::collections::HashMap;
#[tokio::test]
async fn counting_source_resumes_and_ignores_when_non_resumable() {
let s = CountingSource::new(5, 2);
assert_eq!(s.state_key().as_deref(), Some("conformance:counting"));
assert_eq!(s.connector_name(), "counting-source");
assert_eq!(
s.fetch_with_context(&HashMap::new()).await.unwrap().len(),
5
);
s.apply_start_bookmark(json!({ "n": 5 })).await.unwrap();
assert!(
s.fetch_with_context(&HashMap::new())
.await
.unwrap()
.is_empty()
);
let nr = CountingSource::non_resumable(5, 2);
nr.apply_start_bookmark(json!({ "n": 5 })).await.unwrap();
assert_eq!(
nr.fetch_with_context(&HashMap::new()).await.unwrap().len(),
5
);
}
#[tokio::test]
async fn test_sink_accessors() {
let s = TestSink::new();
assert!(s.is_empty());
s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
assert!(!s.is_empty());
assert_eq!(s.len(), 1);
assert_eq!(s.total_written(), 1);
assert_eq!(s.connector_name(), "test-sink");
}
#[tokio::test]
async fn lying_idempotent_sink_never_persists_a_token() {
let s = LyingIdempotentSink::new();
assert!(s.is_empty());
assert!(s.supports_idempotent_writes());
assert_eq!(s.connector_name(), "lying-idempotent-sink");
s.write_batch_idempotent(&[json!({ "id": 1 })], "scope", "00000000000000000001")
.await
.unwrap();
assert_eq!(s.len(), 1);
assert!(s.last_committed_token("scope").await.unwrap().is_none());
}
#[tokio::test]
async fn lying_keyed_sink_appends_duplicates() {
let s = LyingKeyedSink::new();
assert!(s.is_empty());
assert!(s.dedups_by_key());
assert!(s.supported_write_modes().contains(&WriteMode::Upsert));
assert_eq!(s.connector_name(), "lying-keyed-sink");
s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
s.write_batch(&[json!({ "id": 1 })]).await.unwrap();
assert_eq!(s.len(), 2, "lying keyed sink does not dedup");
}
#[tokio::test]
async fn failing_and_panicking_source_labels() {
assert_eq!(FailingSource.connector_name(), "failing-source");
assert_eq!(PanickingSource.connector_name(), "panicking-source");
assert!(FailingSource.fetch_all().await.is_err());
}
#[tokio::test]
async fn test_sink_delete_marker_removes_row() {
let s = TestSink::keyed_upsert("id");
assert!(s.supported_write_modes().contains(&WriteMode::Delete));
s.write_batch(&[json!({ "id": 1, "v": "a" })])
.await
.unwrap();
assert_eq!(s.len(), 1);
s.write_batch(&[json!({ "id": 1, "__op": "d" })])
.await
.unwrap();
assert_eq!(s.len(), 0, "delete marker must remove the row");
let plain = TestSink::keyed("id");
plain
.write_batch(&[json!({ "id": 2, "__op": "d" })])
.await
.unwrap();
assert_eq!(plain.len(), 1, "no marker configured → the row is upserted");
}
#[tokio::test]
async fn evolving_sink_evolves_and_noop_does_not() {
let evo = EvolvingSink::new();
assert_eq!(evo.connector_name(), "evolving-sink");
assert_eq!(evo.write_batch(&[json!({ "id": 1 })]).await.unwrap(), 1);
assert_eq!(evo.column_count(), 1);
let evolution = SchemaEvolution {
additions: vec![ColumnChange {
name: "email".to_string(),
from: None,
to: json!({ "type": "string" }),
}],
widenings: Vec::new(),
relax_nullability: Vec::new(),
};
evo.evolve_schema(&evolution).await.unwrap();
assert_eq!(evo.column_count(), 2);
let schema = evo.current_schema().await.unwrap().unwrap();
assert!(schema["properties"]["email"].is_object());
let noop = NoOpEvolvingSink;
assert!(noop.supports_schema_evolution());
assert_eq!(noop.write_batch(&[json!({ "id": 1 })]).await.unwrap(), 1);
let before = noop.current_schema().await.unwrap().unwrap();
noop.evolve_schema(&evolution).await.unwrap();
let after = noop.current_schema().await.unwrap().unwrap();
assert_eq!(before, after, "noop evolve must not change the schema");
}
#[tokio::test]
async fn multi_page_zero_source_emits_multiple_pages_and_fetches() {
let s = MultiPageZeroSource::new(6);
assert_eq!(s.connector_name(), "multi-page-zero-source");
let ctx: HashMap<String, Value> = HashMap::new();
assert_eq!(s.fetch_with_context(&ctx).await.unwrap().len(), 6);
let mut stream = s.stream_pages(&ctx, 0);
let mut pages = 0usize;
let mut records = 0usize;
while let Some(p) = stream.next().await {
let p = p.unwrap();
pages += 1;
records += p.records.len();
}
assert_eq!(records, 6);
assert!(pages > 1, "must emit more than one page under batch_size=0");
}
#[tokio::test]
async fn empty_name_and_erring_check_doubles() {
assert_eq!(EmptyNameSource.connector_name(), "");
assert!(
EmptyNameSource
.fetch_with_context(&HashMap::new())
.await
.unwrap()
.is_empty()
);
let ctx = CheckContext::default();
assert_eq!(ErringCheckSource.connector_name(), "erring-check-source");
assert!(
ErringCheckSource
.fetch_with_context(&HashMap::new())
.await
.unwrap()
.is_empty()
);
assert!(ErringCheckSource.check(&ctx).await.is_err());
let sink = ErringCheckSink;
assert_eq!(sink.connector_name(), "erring-check-sink");
assert_eq!(sink.write_batch(&[json!({ "x": 1 })]).await.unwrap(), 1);
assert!(sink.check(&ctx).await.is_err());
}
}