use std::{error::Error, future::Future, sync::Arc};
use reifydb::{Database, runtime::context::clock::MockClock};
use reifydb_client::{ChangePayload, SubscriptionConfig, WireFormat, WsClient};
use reifydb_value::value::duration::Duration;
use tokio::{runtime::Runtime, time::timeout};
use crate::common::{
cleanup_server, create_server_instance, create_server_instance_with_clock, start_server_and_get_ws_port,
};
mod basic;
mod batch_mixed_op;
mod data_types;
mod filtered;
mod integration;
mod lifecycle;
mod multiple;
mod notifications;
mod reconnect;
mod shaping;
mod stress;
pub fn unique_table_name(prefix: &str) -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
format!("{}_{}", prefix, timestamp % 1_000_000_000)
}
pub async fn create_test_table(client: &WsClient, name: &str, columns: &[(&str, &str)]) -> Result<(), Box<dyn Error>> {
let _ = client.admin("create namespace test", None).await;
let cols = columns.iter().map(|(name, typ)| format!("{}: {}", name, typ)).collect::<Vec<_>>().join(", ");
client.admin(&format!("create table test::{} {{ {} }}", name, cols), None).await?;
Ok(())
}
pub async fn recv_with_timeout(client: &mut WsClient, timeout_ms: u64) -> Option<ChangePayload> {
match timeout(Duration::from_milliseconds(timeout_ms as i64).unwrap().to_std(), client.recv()).await {
Ok(result) => result,
Err(_) => None,
}
}
pub async fn recv_multiple_with_timeout(client: &mut WsClient, count: usize, timeout_ms: u64) -> Vec<ChangePayload> {
let mut results = Vec::new();
let deadline = tokio::time::Instant::now() + Duration::from_milliseconds(timeout_ms as i64).unwrap().to_std();
while results.len() < count {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
break;
}
match timeout(remaining, client.recv()).await {
Ok(Some(change)) => results.push(change),
Ok(None) => break,
Err(_) => break,
}
}
results
}
pub struct JsonColumn {
pub name: String,
pub r#type: String,
pub payload: Vec<String>,
}
pub fn extract_columns(body: &serde_json::Value) -> Vec<JsonColumn> {
let frames = body.get("frames").and_then(|f| f.as_array());
let frame = frames.and_then(|f| f.first());
let columns = frame.and_then(|f| f.get("columns")).and_then(|c| c.as_array());
match columns {
Some(cols) => cols
.iter()
.map(|c| JsonColumn {
name: c.get("name").and_then(|n| n.as_str()).unwrap_or("").to_string(),
r#type: c.get("type").and_then(|t| t.as_str()).unwrap_or("").to_string(),
payload: c
.get("payload")
.and_then(|d| d.as_array())
.map(|arr| {
arr.iter()
.map(|v| match v {
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Number(n) => n.to_string(),
serde_json::Value::Bool(b) => b.to_string(),
serde_json::Value::Null => "null".to_string(),
other => other.to_string(),
})
.collect()
})
.unwrap_or_default(),
})
.collect(),
None => Vec::new(),
}
}
pub fn find_column(body: &serde_json::Value, name: &str) -> Option<JsonColumn> {
extract_columns(body).into_iter().find(|c| c.name == name)
}
pub fn get_op_value(body: &serde_json::Value, row_index: usize) -> Option<i32> {
find_column(body, "_op").and_then(|col| col.payload.get(row_index).cloned()).and_then(|s| s.parse::<i32>().ok())
}
pub struct SubscriptionTestHarness;
impl SubscriptionTestHarness {
pub fn run<F, Fut>(test_fn: F)
where
F: Fn(TestContext) -> Fut + Send + Sync,
Fut: Future<Output = Result<(), Box<dyn Error>>>,
{
Self::run_on(create_server_instance, test_fn)
}
#[allow(dead_code)]
pub fn run_with_clock<F, Fut>(test_fn: F)
where
F: Fn(TestContext, MockClock) -> Fut + Send + Sync,
Fut: Future<Output = Result<(), Box<dyn Error>>>,
{
let runtime = Arc::new(Runtime::new().unwrap());
let _guard = runtime.enter();
let clock = MockClock::from_millis(0);
let mut server = create_server_instance_with_clock(&runtime, clock.clone());
let port = start_server_and_get_ws_port(&runtime, &mut server).unwrap();
runtime.block_on(async {
let mut client =
WsClient::connect(&format!("ws://[::1]:{}", port), WireFormat::Frames).await.unwrap();
client.authenticate("mysecrettoken").await.unwrap();
let ctx = TestContext::new(client);
test_fn(ctx, clock).await.unwrap();
});
cleanup_server(Some(server));
}
fn run_on<F, Fut>(make_server: fn(&Arc<Runtime>) -> Database, test_fn: F)
where
F: Fn(TestContext) -> Fut + Send + Sync,
Fut: Future<Output = Result<(), Box<dyn Error>>>,
{
let runtime = Arc::new(Runtime::new().unwrap());
let _guard = runtime.enter();
let mut server = make_server(&runtime);
let port = start_server_and_get_ws_port(&runtime, &mut server).unwrap();
runtime.block_on(async {
let mut client =
WsClient::connect(&format!("ws://[::1]:{}", port), WireFormat::Frames).await.unwrap();
client.authenticate("mysecrettoken").await.unwrap();
let ctx = TestContext::new(client);
test_fn(ctx).await.unwrap();
});
cleanup_server(Some(server));
}
}
pub struct TestContext {
pub client: WsClient,
table_prefix: String,
}
impl TestContext {
fn new(client: WsClient) -> Self {
Self {
client,
table_prefix: unique_table_name("t"),
}
}
pub async fn rql(&self, query: &str) -> Result<(), Box<dyn Error>> {
self.client.command(query, None).await?;
Ok(())
}
pub async fn create_table(&self, name: &str, columns: &str) -> Result<String, Box<dyn Error>> {
let full_name = format!("{}_{}", self.table_prefix, name);
let _ = self.client.admin("create namespace test", None).await;
self.client.admin(&format!("create table test::{} {{ {} }}", full_name, columns), None).await?;
Ok(full_name)
}
pub async fn subscribe(&mut self, table: &str, config: SubscriptionConfig) -> Result<String, Box<dyn Error>> {
let sub_id = self.client.subscribe(&format!("from test::{}", table), config).await?;
Ok(sub_id)
}
pub async fn insert(&self, table: &str, rows: &str) -> Result<(), Box<dyn Error>> {
self.client.command(&format!("INSERT test::{} [{}]", table, rows), None).await?;
Ok(())
}
pub async fn update(&self, table: &str, filter: &str, map: &str) -> Result<(), Box<dyn Error>> {
self.client
.command(&format!("UPDATE test::{} {{ {} }} FILTER {{{}}}", table, map, filter), None)
.await?;
Ok(())
}
pub async fn delete(&self, table: &str, filter: &str) -> Result<(), Box<dyn Error>> {
self.client.command(&format!("DELETE test::{} FILTER {{{}}}", table, filter), None).await?;
Ok(())
}
pub async fn recv(&mut self) -> Option<ChangePayload> {
recv_with_timeout(&mut self.client, 5000).await
}
pub async fn close(self, sub_id: &str) -> Result<(), Box<dyn Error>> {
self.client.unsubscribe(sub_id).await?;
self.client.close().await?;
Ok(())
}
}