496
crates/sq-server/tests/data_plane_test.rs
Normal file
496
crates/sq-server/tests/data_plane_test.rs
Normal file
@@ -0,0 +1,496 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use sq_grpc_interface::{
|
||||
data_plane_service_client::DataPlaneServiceClient,
|
||||
data_plane_service_server::DataPlaneServiceServer,
|
||||
status_service_client::StatusServiceClient,
|
||||
status_service_server::StatusServiceServer,
|
||||
AckMode, GetStatusRequest, MessageHeader, PublishMessage, PublishRequest, PublishSettings,
|
||||
SubscribeRequest,
|
||||
};
|
||||
use sq_sim::fs::InMemoryFileSystem;
|
||||
use sq_sim::SimClock;
|
||||
use sq_storage::engine::StorageEngine;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_stream::StreamExt;
|
||||
|
||||
/// A lightweight test harness that starts a gRPC server on a random port
|
||||
/// and returns both the server task and connected clients.
|
||||
struct TestServer {
|
||||
addr: SocketAddr,
|
||||
_shutdown: tokio::sync::oneshot::Sender<()>,
|
||||
}
|
||||
|
||||
impl TestServer {
|
||||
async fn start() -> Self {
|
||||
let fs = Arc::new(InMemoryFileSystem::new());
|
||||
let clock = Arc::new(SimClock::new());
|
||||
let config = sq_models::WalConfig {
|
||||
max_segment_bytes: 1024 * 1024,
|
||||
max_segment_age_secs: 3600,
|
||||
data_dir: PathBuf::from("/data"),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let engine = StorageEngine::new(fs, clock, config).unwrap();
|
||||
engine.recover().unwrap();
|
||||
|
||||
let engine = Arc::new(Mutex::new(engine));
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
|
||||
// Build the health server state-like object inline for tests.
|
||||
let node_id = "test-node".to_string();
|
||||
|
||||
struct TestHealthServer {
|
||||
node_id: String,
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl sq_grpc_interface::status_service_server::StatusService for TestHealthServer {
|
||||
async fn status(
|
||||
&self,
|
||||
_request: tonic::Request<GetStatusRequest>,
|
||||
) -> Result<tonic::Response<sq_grpc_interface::GetStatusResponse>, tonic::Status> {
|
||||
Ok(tonic::Response::new(sq_grpc_interface::GetStatusResponse {
|
||||
node_id: self.node_id.clone(),
|
||||
cluster: None,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
struct TestDataPlaneServer {
|
||||
engine: Arc<Mutex<StorageEngine<InMemoryFileSystem, SimClock>>>,
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl sq_grpc_interface::data_plane_service_server::DataPlaneService
|
||||
for TestDataPlaneServer
|
||||
{
|
||||
async fn publish(
|
||||
&self,
|
||||
request: tonic::Request<PublishRequest>,
|
||||
) -> Result<tonic::Response<sq_grpc_interface::PublishResponse>, tonic::Status> {
|
||||
let req = request.into_inner();
|
||||
|
||||
if req.messages.is_empty() {
|
||||
return Err(tonic::Status::invalid_argument(
|
||||
"messages must not be empty",
|
||||
));
|
||||
}
|
||||
|
||||
let mut results = Vec::new();
|
||||
let engine = self.engine.lock().await;
|
||||
|
||||
for msg in &req.messages {
|
||||
if msg.topic.is_empty() {
|
||||
return Err(tonic::Status::invalid_argument("topic must not be empty"));
|
||||
}
|
||||
|
||||
let headers: Vec<sq_models::Header> = msg
|
||||
.headers
|
||||
.iter()
|
||||
.map(|h| sq_models::Header {
|
||||
key: h.key.clone(),
|
||||
value: h.value.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let key = if msg.key.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(msg.key.as_slice())
|
||||
};
|
||||
|
||||
let offset = engine
|
||||
.append(&msg.topic, 0, key, &msg.value, &headers, 0)
|
||||
.map_err(|e| tonic::Status::internal(e.to_string()))?;
|
||||
|
||||
results.push(sq_grpc_interface::PublishResult {
|
||||
topic: msg.topic.clone(),
|
||||
partition: 0,
|
||||
offset,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tonic::Response::new(sq_grpc_interface::PublishResponse {
|
||||
results,
|
||||
}))
|
||||
}
|
||||
|
||||
type SubscribeStream = std::pin::Pin<
|
||||
Box<
|
||||
dyn tokio_stream::Stream<
|
||||
Item = Result<sq_grpc_interface::SubscribeResponse, tonic::Status>,
|
||||
> + Send
|
||||
+ 'static,
|
||||
>,
|
||||
>;
|
||||
|
||||
async fn subscribe(
|
||||
&self,
|
||||
request: tonic::Request<SubscribeRequest>,
|
||||
) -> Result<tonic::Response<Self::SubscribeStream>, tonic::Status> {
|
||||
let req = request.into_inner();
|
||||
let batch_size = if req.max_batch_size == 0 {
|
||||
100
|
||||
} else {
|
||||
req.max_batch_size as usize
|
||||
};
|
||||
let start_offset = req.start_offset.unwrap_or(0);
|
||||
let topic = req.topic.clone();
|
||||
let partition = req.partition;
|
||||
let engine = self.engine.clone();
|
||||
|
||||
let stream = async_stream::try_stream! {
|
||||
let mut current_offset = start_offset;
|
||||
let mut empty_polls = 0u32;
|
||||
|
||||
loop {
|
||||
let messages = {
|
||||
let eng = engine.lock().await;
|
||||
eng.read(&topic, partition, current_offset, batch_size)
|
||||
.map_err(|e| tonic::Status::internal(e.to_string()))?
|
||||
};
|
||||
|
||||
if messages.is_empty() {
|
||||
empty_polls += 1;
|
||||
// In tests, stop after a few empty polls to avoid hanging.
|
||||
if empty_polls > 3 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
empty_polls = 0;
|
||||
|
||||
let consumed: Vec<sq_grpc_interface::ConsumedMessage> = messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
current_offset = m.offset + 1;
|
||||
sq_grpc_interface::ConsumedMessage {
|
||||
offset: m.offset,
|
||||
topic: m.topic.to_string(),
|
||||
partition: m.partition,
|
||||
key: m.key.clone().unwrap_or_default(),
|
||||
value: m.value.clone(),
|
||||
headers: m
|
||||
.headers
|
||||
.iter()
|
||||
.map(|h| MessageHeader {
|
||||
key: h.key.clone(),
|
||||
value: h.value.clone(),
|
||||
})
|
||||
.collect(),
|
||||
timestamp_ms: m.timestamp_ms,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
yield sq_grpc_interface::SubscribeResponse { messages: consumed };
|
||||
}
|
||||
};
|
||||
|
||||
Ok(tonic::Response::new(Box::pin(stream)))
|
||||
}
|
||||
|
||||
async fn ack(
|
||||
&self,
|
||||
_request: tonic::Request<sq_grpc_interface::AckRequest>,
|
||||
) -> Result<tonic::Response<sq_grpc_interface::AckResponse>, tonic::Status> {
|
||||
Ok(tonic::Response::new(sq_grpc_interface::AckResponse {}))
|
||||
}
|
||||
}
|
||||
|
||||
let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
|
||||
|
||||
tokio::spawn(async move {
|
||||
tonic::transport::Server::builder()
|
||||
.add_service(StatusServiceServer::new(TestHealthServer {
|
||||
node_id: node_id.clone(),
|
||||
}))
|
||||
.add_service(DataPlaneServiceServer::new(TestDataPlaneServer {
|
||||
engine,
|
||||
}))
|
||||
.serve_with_incoming_shutdown(incoming, async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
// Give the server a moment to start.
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||
|
||||
TestServer {
|
||||
addr,
|
||||
_shutdown: shutdown_tx,
|
||||
}
|
||||
}
|
||||
|
||||
fn endpoint(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_status_rpc() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = StatusServiceClient::connect(server.endpoint()).await.unwrap();
|
||||
|
||||
let response = client
|
||||
.status(tonic::Request::new(GetStatusRequest {}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.into_inner().node_id, "test-node");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_publish_single_message() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages: vec![PublishMessage {
|
||||
topic: "orders".to_string(),
|
||||
key: vec![],
|
||||
value: b"hello world".to_vec(),
|
||||
headers: vec![],
|
||||
}],
|
||||
settings: None,
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let results = response.into_inner().results;
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0].topic, "orders");
|
||||
assert_eq!(results[0].offset, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_publish_batch_sequential_offsets() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let messages: Vec<PublishMessage> = (0..100)
|
||||
.map(|i| PublishMessage {
|
||||
topic: "events".to_string(),
|
||||
key: vec![],
|
||||
value: format!("msg-{i}").into_bytes(),
|
||||
headers: vec![],
|
||||
})
|
||||
.collect();
|
||||
|
||||
let response = client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages,
|
||||
settings: Some(PublishSettings {
|
||||
ack_mode: AckMode::All.into(),
|
||||
}),
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let results = response.into_inner().results;
|
||||
assert_eq!(results.len(), 100);
|
||||
for (i, r) in results.iter().enumerate() {
|
||||
assert_eq!(r.offset, i as u64);
|
||||
assert_eq!(r.topic, "events");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_publish_empty_topic_returns_error() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let err = client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages: vec![PublishMessage {
|
||||
topic: "".to_string(),
|
||||
key: vec![],
|
||||
value: b"data".to_vec(),
|
||||
headers: vec![],
|
||||
}],
|
||||
settings: None,
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(err.code(), tonic::Code::InvalidArgument);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_publish_empty_messages_returns_error() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let err = client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages: vec![],
|
||||
settings: None,
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(err.code(), tonic::Code::InvalidArgument);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_publish_with_key_and_headers() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages: vec![PublishMessage {
|
||||
topic: "orders".to_string(),
|
||||
key: b"order-123".to_vec(),
|
||||
value: b"payload".to_vec(),
|
||||
headers: vec![MessageHeader {
|
||||
key: "trace-id".to_string(),
|
||||
value: b"abc-123".to_vec(),
|
||||
}],
|
||||
}],
|
||||
settings: None,
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let results = response.into_inner().results;
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0].offset, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_subscribe_from_beginning() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Publish 10 messages first.
|
||||
let messages: Vec<PublishMessage> = (0..10)
|
||||
.map(|i| PublishMessage {
|
||||
topic: "events".to_string(),
|
||||
key: vec![],
|
||||
value: format!("msg-{i}").into_bytes(),
|
||||
headers: vec![],
|
||||
})
|
||||
.collect();
|
||||
|
||||
client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages,
|
||||
settings: None,
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Subscribe from offset 0.
|
||||
let response = client
|
||||
.subscribe(tonic::Request::new(SubscribeRequest {
|
||||
topic: "events".to_string(),
|
||||
partition: 0,
|
||||
consumer_group: "".to_string(),
|
||||
start_offset: Some(0),
|
||||
max_batch_size: 100,
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut stream = response.into_inner();
|
||||
let mut all_messages = Vec::new();
|
||||
|
||||
while let Some(Ok(batch)) = stream.next().await {
|
||||
all_messages.extend(batch.messages);
|
||||
if all_messages.len() >= 10 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(all_messages.len(), 10);
|
||||
for (i, msg) in all_messages.iter().enumerate() {
|
||||
assert_eq!(msg.offset, i as u64);
|
||||
assert_eq!(msg.value, format!("msg-{i}").as_bytes());
|
||||
assert_eq!(msg.topic, "events");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_subscribe_from_middle() {
|
||||
let server = TestServer::start().await;
|
||||
let mut client = DataPlaneServiceClient::connect(server.endpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Publish 10 messages.
|
||||
let messages: Vec<PublishMessage> = (0..10)
|
||||
.map(|i| PublishMessage {
|
||||
topic: "events".to_string(),
|
||||
key: vec![],
|
||||
value: format!("msg-{i}").into_bytes(),
|
||||
headers: vec![],
|
||||
})
|
||||
.collect();
|
||||
|
||||
client
|
||||
.publish(tonic::Request::new(PublishRequest {
|
||||
messages,
|
||||
settings: None,
|
||||
producer_id: "test".to_string(),
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Subscribe from offset 5.
|
||||
let response = client
|
||||
.subscribe(tonic::Request::new(SubscribeRequest {
|
||||
topic: "events".to_string(),
|
||||
partition: 0,
|
||||
consumer_group: "".to_string(),
|
||||
start_offset: Some(5),
|
||||
max_batch_size: 100,
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut stream = response.into_inner();
|
||||
let mut all_messages = Vec::new();
|
||||
|
||||
while let Some(Ok(batch)) = stream.next().await {
|
||||
all_messages.extend(batch.messages);
|
||||
if all_messages.len() >= 5 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(all_messages.len(), 5);
|
||||
assert_eq!(all_messages[0].offset, 5);
|
||||
assert_eq!(all_messages[4].offset, 9);
|
||||
}
|
||||
Reference in New Issue
Block a user