328 lines
11 KiB
Rust
328 lines
11 KiB
Rust
use redis::{Client, Commands, Connection, RedisResult};
|
|
use std::process::{Child, Command};
|
|
use std::time::Duration;
|
|
use tokio::time::sleep;
|
|
|
|
// Helper function to get Redis connection, retrying until successful
|
|
fn get_redis_connection(port: u16) -> Connection {
|
|
let connection_info = format!("redis://127.0.0.1:{}", port);
|
|
let client = Client::open(connection_info).unwrap();
|
|
let mut attempts = 0;
|
|
loop {
|
|
match client.get_connection() {
|
|
Ok(mut conn) => {
|
|
if redis::cmd("PING").query::<String>(&mut conn).is_ok() {
|
|
// Acquire ReadWrite permissions on this connection
|
|
let sel: RedisResult<String> = redis::cmd("SELECT")
|
|
.arg(0)
|
|
.arg("KEY")
|
|
.arg("test-admin")
|
|
.query(&mut conn);
|
|
if sel.is_ok() {
|
|
return conn;
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
if attempts >= 120 {
|
|
panic!(
|
|
"Failed to connect to Redis server after 120 attempts: {}",
|
|
e
|
|
);
|
|
}
|
|
}
|
|
}
|
|
attempts += 1;
|
|
std::thread::sleep(Duration::from_millis(100));
|
|
}
|
|
}
|
|
|
|
// A guard to ensure the server process is killed when it goes out of scope
|
|
struct ServerProcessGuard {
|
|
process: Child,
|
|
test_dir: String,
|
|
}
|
|
|
|
impl Drop for ServerProcessGuard {
|
|
fn drop(&mut self) {
|
|
println!("Killing server process (pid: {})...", self.process.id());
|
|
if let Err(e) = self.process.kill() {
|
|
eprintln!("Failed to kill server process: {}", e);
|
|
}
|
|
match self.process.wait() {
|
|
Ok(status) => println!("Server process exited with: {}", status),
|
|
Err(e) => eprintln!("Failed to wait on server process: {}", e),
|
|
}
|
|
|
|
// Clean up the specific test directory
|
|
println!("Cleaning up test directory: {}", self.test_dir);
|
|
if let Err(e) = std::fs::remove_dir_all(&self.test_dir) {
|
|
eprintln!("Failed to clean up test directory: {}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Helper to set up the server and return a connection
|
|
fn setup_server() -> (ServerProcessGuard, u16) {
|
|
use std::sync::atomic::{AtomicU16, Ordering};
|
|
static PORT_COUNTER: AtomicU16 = AtomicU16::new(16400);
|
|
let port = PORT_COUNTER.fetch_add(1, Ordering::SeqCst);
|
|
|
|
let test_dir = format!("/tmp/herodb_test_{}", port);
|
|
|
|
// Clean up previous test data
|
|
if std::path::Path::new(&test_dir).exists() {
|
|
let _ = std::fs::remove_dir_all(&test_dir);
|
|
}
|
|
std::fs::create_dir_all(&test_dir).unwrap();
|
|
|
|
// Start the server in a subprocess
|
|
let child = Command::new("cargo")
|
|
.args(&[
|
|
"run",
|
|
"--",
|
|
"--dir",
|
|
&test_dir,
|
|
"--port",
|
|
&port.to_string(),
|
|
"--debug",
|
|
"--admin-secret",
|
|
"test-admin",
|
|
])
|
|
.spawn()
|
|
.expect("Failed to start server process");
|
|
|
|
// Create a new guard that also owns the test directory path
|
|
let guard = ServerProcessGuard {
|
|
process: child,
|
|
test_dir,
|
|
};
|
|
|
|
// Give the server time to build and start (cargo run may compile first)
|
|
std::thread::sleep(Duration::from_millis(2500));
|
|
|
|
(guard, port)
|
|
}
|
|
|
|
async fn cleanup_keys(conn: &mut Connection) {
|
|
let keys: Vec<String> = redis::cmd("KEYS").arg("*").query(conn).unwrap();
|
|
if !keys.is_empty() {
|
|
for key in keys {
|
|
let _: () = redis::cmd("DEL").arg(key).query(conn).unwrap();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn all_tests() {
|
|
let (_server_guard, port) = setup_server();
|
|
let mut conn = get_redis_connection(port);
|
|
|
|
// Run all tests using the same connection
|
|
test_basic_ping(&mut conn).await;
|
|
test_string_operations(&mut conn).await;
|
|
test_incr_operations(&mut conn).await;
|
|
test_hash_operations(&mut conn).await;
|
|
test_expiration(&mut conn).await;
|
|
test_scan_operations(&mut conn).await;
|
|
test_scan_with_count(&mut conn).await;
|
|
test_hscan_operations(&mut conn).await;
|
|
test_transaction_operations(&mut conn).await;
|
|
test_discard_transaction(&mut conn).await;
|
|
test_type_command(&mut conn).await;
|
|
test_info_command(&mut conn).await;
|
|
}
|
|
|
|
async fn test_basic_ping(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let result: String = redis::cmd("PING").query(conn).unwrap();
|
|
assert_eq!(result, "PONG");
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_string_operations(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let _: () = conn.set("key", "value").unwrap();
|
|
let result: String = conn.get("key").unwrap();
|
|
assert_eq!(result, "value");
|
|
let result: Option<String> = conn.get("noexist").unwrap();
|
|
assert_eq!(result, None);
|
|
let deleted: i32 = conn.del("key").unwrap();
|
|
assert_eq!(deleted, 1);
|
|
let result: Option<String> = conn.get("key").unwrap();
|
|
assert_eq!(result, None);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_incr_operations(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let result: i32 = redis::cmd("INCR").arg("counter").query(conn).unwrap();
|
|
assert_eq!(result, 1);
|
|
let result: i32 = redis::cmd("INCR").arg("counter").query(conn).unwrap();
|
|
assert_eq!(result, 2);
|
|
let _: () = conn.set("string", "hello").unwrap();
|
|
let result: RedisResult<i32> = redis::cmd("INCR").arg("string").query(conn);
|
|
assert!(result.is_err());
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_hash_operations(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let result: i32 = conn.hset("hash", "field1", "value1").unwrap();
|
|
assert_eq!(result, 1);
|
|
let result: String = conn.hget("hash", "field1").unwrap();
|
|
assert_eq!(result, "value1");
|
|
let _: () = conn.hset("hash", "field2", "value2").unwrap();
|
|
let _: () = conn.hset("hash", "field3", "value3").unwrap();
|
|
let result: std::collections::HashMap<String, String> = conn.hgetall("hash").unwrap();
|
|
assert_eq!(result.len(), 3);
|
|
assert_eq!(result.get("field1").unwrap(), "value1");
|
|
assert_eq!(result.get("field2").unwrap(), "value2");
|
|
assert_eq!(result.get("field3").unwrap(), "value3");
|
|
let result: i32 = conn.hlen("hash").unwrap();
|
|
assert_eq!(result, 3);
|
|
let result: bool = conn.hexists("hash", "field1").unwrap();
|
|
assert_eq!(result, true);
|
|
let result: bool = conn.hexists("hash", "noexist").unwrap();
|
|
assert_eq!(result, false);
|
|
let result: i32 = conn.hdel("hash", "field1").unwrap();
|
|
assert_eq!(result, 1);
|
|
let mut result: Vec<String> = conn.hkeys("hash").unwrap();
|
|
result.sort();
|
|
assert_eq!(result, vec!["field2", "field3"]);
|
|
let mut result: Vec<String> = conn.hvals("hash").unwrap();
|
|
result.sort();
|
|
assert_eq!(result, vec!["value2", "value3"]);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_expiration(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let _: () = conn.set_ex("expkey", "value", 1).unwrap();
|
|
let result: i32 = conn.ttl("expkey").unwrap();
|
|
assert!(result == 1 || result == 0);
|
|
let result: bool = conn.exists("expkey").unwrap();
|
|
assert_eq!(result, true);
|
|
sleep(Duration::from_millis(1100)).await;
|
|
let result: Option<String> = conn.get("expkey").unwrap();
|
|
assert_eq!(result, None);
|
|
let result: i32 = conn.ttl("expkey").unwrap();
|
|
assert_eq!(result, -2);
|
|
let result: bool = conn.exists("expkey").unwrap();
|
|
assert_eq!(result, false);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_scan_operations(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
for i in 0..5 {
|
|
let _: () = conn.set(format!("key{}", i), format!("value{}", i)).unwrap();
|
|
}
|
|
let result: (u64, Vec<String>) = redis::cmd("SCAN")
|
|
.arg(0)
|
|
.arg("MATCH")
|
|
.arg("key*")
|
|
.arg("COUNT")
|
|
.arg(10)
|
|
.query(conn)
|
|
.unwrap();
|
|
let (cursor, keys) = result;
|
|
assert_eq!(cursor, 0);
|
|
assert_eq!(keys.len(), 5);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_scan_with_count(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
for i in 0..15 {
|
|
let _: () = conn.set(format!("scan_key{}", i), i).unwrap();
|
|
}
|
|
let mut cursor = 0;
|
|
let mut all_keys = std::collections::HashSet::new();
|
|
loop {
|
|
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
|
|
.arg(cursor)
|
|
.arg("MATCH")
|
|
.arg("scan_key*")
|
|
.arg("COUNT")
|
|
.arg(5)
|
|
.query(conn)
|
|
.unwrap();
|
|
for key in keys {
|
|
all_keys.insert(key);
|
|
}
|
|
cursor = next_cursor;
|
|
if cursor == 0 {
|
|
break;
|
|
}
|
|
}
|
|
assert_eq!(all_keys.len(), 15);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_hscan_operations(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
for i in 0..3 {
|
|
let _: () = conn.hset("testhash", format!("field{}", i), format!("value{}", i)).unwrap();
|
|
}
|
|
let result: (u64, Vec<String>) = redis::cmd("HSCAN")
|
|
.arg("testhash")
|
|
.arg(0)
|
|
.arg("MATCH")
|
|
.arg("*")
|
|
.arg("COUNT")
|
|
.arg(10)
|
|
.query(conn)
|
|
.unwrap();
|
|
let (cursor, fields) = result;
|
|
assert_eq!(cursor, 0);
|
|
assert_eq!(fields.len(), 6);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_transaction_operations(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let _: () = redis::cmd("MULTI").query(conn).unwrap();
|
|
let _: () = redis::cmd("SET").arg("key1").arg("value1").query(conn).unwrap();
|
|
let _: () = redis::cmd("SET").arg("key2").arg("value2").query(conn).unwrap();
|
|
let _: Vec<String> = redis::cmd("EXEC").query(conn).unwrap();
|
|
let result: String = conn.get("key1").unwrap();
|
|
assert_eq!(result, "value1");
|
|
let result: String = conn.get("key2").unwrap();
|
|
assert_eq!(result, "value2");
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_discard_transaction(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let _: () = redis::cmd("MULTI").query(conn).unwrap();
|
|
let _: () = redis::cmd("SET").arg("discard").arg("value").query(conn).unwrap();
|
|
let _: () = redis::cmd("DISCARD").query(conn).unwrap();
|
|
let result: Option<String> = conn.get("discard").unwrap();
|
|
assert_eq!(result, None);
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
async fn test_type_command(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let _: () = conn.set("string", "value").unwrap();
|
|
let result: String = redis::cmd("TYPE").arg("string").query(conn).unwrap();
|
|
assert_eq!(result, "string");
|
|
let _: () = conn.hset("hash", "field", "value").unwrap();
|
|
let result: String = redis::cmd("TYPE").arg("hash").query(conn).unwrap();
|
|
assert_eq!(result, "hash");
|
|
let result: String = redis::cmd("TYPE").arg("noexist").query(conn).unwrap();
|
|
assert_eq!(result, "none");
|
|
cleanup_keys(conn).await;
|
|
}
|
|
|
|
|
|
async fn test_info_command(conn: &mut Connection) {
|
|
cleanup_keys(conn).await;
|
|
let result: String = redis::cmd("INFO").query(conn).unwrap();
|
|
assert!(result.contains("redis_version"));
|
|
let result: String = redis::cmd("INFO").arg("replication").query(conn).unwrap();
|
|
assert!(result.contains("role:master"));
|
|
cleanup_keys(conn).await;
|
|
}
|