use processor::kafka::{CommitStrategy, KafkaSource, KafkaSourceConfig};
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct TestEvent {
id: String,
value: f64,
timestamp: i64,
}
fn create_test_config(group_id: &str, topic: &str) -> KafkaSourceConfig {
KafkaSourceConfig {
brokers: "localhost:9092".to_string(),
group_id: group_id.to_string(),
topics: vec![topic.to_string()],
commit_strategy: CommitStrategy::Periodic,
commit_interval_secs: Some(1),
auto_offset_reset: Some("earliest".to_string()),
poll_timeout_ms: Some(1000),
..Default::default()
}
}
#[tokio::test]
#[ignore] async fn test_kafka_source_creation() {
let config = create_test_config("test-group-1", "test-topic-1");
let result = KafkaSource::<TestEvent>::new(config, None).await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore] async fn test_kafka_source_subscribe() {
let config = create_test_config("test-group-2", "test-topic-2");
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let result = source.subscribe(&["test-topic-2".to_string()]).await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore] async fn test_kafka_source_poll_timeout() {
let config = KafkaSourceConfig {
poll_timeout_ms: Some(100),
..create_test_config("test-group-3", "empty-topic")
};
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
source.subscribe(&["empty-topic".to_string()]).await.unwrap();
let result = source.poll().await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore] async fn test_kafka_source_metrics() {
let config = create_test_config("test-group-4", "test-topic-4");
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let metrics = source.metrics().await;
assert_eq!(metrics.messages_consumed, 0);
assert_eq!(metrics.messages_failed, 0);
assert_eq!(metrics.total_commits, 0);
}
#[tokio::test]
#[ignore] async fn test_kafka_source_start_and_stop() {
let config = create_test_config("test-group-5", "test-topic-5");
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let (tx, mut rx) = mpsc::channel(10);
let source_clone = std::sync::Arc::new(source);
let source_for_task = source_clone.clone();
let consumer_task = tokio::spawn(async move {
source_for_task.start(tx).await
});
tokio::time::sleep(Duration::from_secs(2)).await;
source_clone.close().await.expect("Failed to close source");
let result = tokio::time::timeout(Duration::from_secs(5), consumer_task).await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore] async fn test_kafka_source_commit_strategies() {
let config = KafkaSourceConfig {
commit_strategy: CommitStrategy::Manual,
..create_test_config("test-group-6", "test-topic-6")
};
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let result = source.commit().await;
assert!(result.is_ok());
let config = KafkaSourceConfig {
commit_strategy: CommitStrategy::Periodic,
commit_interval_secs: Some(1),
..create_test_config("test-group-7", "test-topic-7")
};
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let result = source.commit().await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore] async fn test_kafka_source_multiple_topics() {
let config = KafkaSourceConfig {
topics: vec!["topic-a".to_string(), "topic-b".to_string()],
..create_test_config("test-group-8", "topic-a")
};
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let result = source
.subscribe(&["topic-a".to_string(), "topic-b".to_string()])
.await;
assert!(result.is_ok());
}
#[tokio::test]
#[ignore] async fn test_kafka_source_message_consumption() {
let config = create_test_config("test-group-9", "test-messages");
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
let (tx, mut rx) = mpsc::channel(100);
let source = std::sync::Arc::new(source);
let source_for_task = source.clone();
let consumer_task = tokio::spawn(async move {
source_for_task.start(tx).await
});
let timeout = tokio::time::sleep(Duration::from_secs(5));
tokio::pin!(timeout);
let mut count = 0;
loop {
tokio::select! {
Some(msg) = rx.recv() => {
count += 1;
assert!(msg.offset >= 0);
assert!(!msg.topic.is_empty());
println!("Consumed message: {:?}", msg.payload);
if count >= 10 {
break;
}
}
_ = &mut timeout => {
break;
}
}
}
println!("Consumed {} messages", count);
source.close().await.expect("Failed to close");
let metrics = source.metrics().await;
assert_eq!(metrics.messages_consumed, count);
}
#[tokio::test]
#[ignore] async fn test_kafka_source_lag_tracking() {
let config = create_test_config("test-group-10", "test-lag");
let source = KafkaSource::<TestEvent>::new(config, None)
.await
.expect("Failed to create source");
source.subscribe(&["test-lag".to_string()]).await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
let result = source.get_lag().await;
if let Ok(lag) = result {
println!("Lag information: {:?}", lag);
}
}