Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use OS-picked ports in the cli.rs tests #150

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 79 additions & 26 deletions chitchat-test/tests/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@

mod helpers;

use std::net::TcpListener;
use std::process::Child;
use std::thread;
use std::time::Duration;
use std::{thread, time::Duration};

use chitchat_test::{ApiResponse, SetKeyValueResponse};
use helpers::spawn_command;
Expand All @@ -17,33 +17,82 @@ impl Drop for KillOnDrop {
}
}

struct NodeProcess {
node_id: String,
port: u16,
seed_addr: Option<String>,
_child: Option<KillOnDrop>, // Never accessed, but included to keep the value alive.
}

impl NodeProcess {
pub fn new(node_id: String, port: u16, seed_addr: Option<&str>) -> Self {
Self {
node_id,
port,
seed_addr: seed_addr.map(|x| x.to_string()),
_child: None,
}
}

fn start(&mut self) -> anyhow::Result<()> {
// Construct the command arguments
let mut args = vec![
"--listen_addr".to_string(),
format!("127.0.0.1:{}", self.port),
"--node_id".to_string(),
self.node_id.clone(),
"--interval_ms".to_string(),
"50".to_string(),
];

if let Some(seed) = &self.seed_addr {
args.push("--seed".to_string());
args.push(seed.clone());
}

// Start the node process
let child = spawn_command(&args.join(" "))?;
self._child = Some(KillOnDrop(child));
return Ok(());
}

fn base_url(&self) -> String {
format!("http://127.0.0.1:{}", self.port)
}
}

fn create_node_process(node_id: &str, seed_addr: Option<&str>) -> anyhow::Result<NodeProcess> {
// First, find a free port by binding to port 0 to let the OS choose an available port.
let listener = TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);

// Second, use the port to construct a new node process.
let mut process = NodeProcess::new(node_id.to_string(), port, seed_addr);
process.start()?;

Ok(process)
}

fn setup_nodes(
port_offset: usize,
num_nodes: usize,
wait_stabilization_secs: u64,
seed_requires_dns: bool,
) -> Vec<KillOnDrop> {
let seed_port = port_offset;
let command_args =
format!("--listen_addr 127.0.0.1:{seed_port} --node_id node_0 --interval_ms 120");
let seed_node = spawn_command(&command_args).unwrap();
) -> Vec<NodeProcess> {
let child_handle = create_node_process("node_0", None).unwrap();
let seed_port = child_handle.port;
let mut child_process_handles = vec![child_handle];

let mut child_process_handles = vec![KillOnDrop(seed_node)];
let seed_host_name = if seed_requires_dns {
"localhost"
} else {
"127.0.0.1"
};
let seed_addr = format!("{}:{}", seed_host_name, seed_port);

for i in 1..num_nodes {
let node_id = format!("node_{i}");
let node_port = port_offset + i;
let seed_host_name = if seed_requires_dns {
"localhost"
} else {
"127.0.0.1"
};
let command_args = format!(
"--listen_addr 127.0.0.1:{node_port} --seed {seed_host_name}:{seed_port} --node_id \
{node_id} --interval_ms 50"
);
let node = spawn_command(&command_args).unwrap();
child_process_handles.push(KillOnDrop(node));
child_process_handles.push(create_node_process(&node_id, Some(&seed_addr)).unwrap());
}
thread::sleep(Duration::from_secs(wait_stabilization_secs));
child_process_handles
Expand All @@ -61,16 +110,18 @@ fn set_kv(node_api_endpoint: &str, key: &str, value: &str) -> anyhow::Result<Set
}
#[test]
fn test_multiple_nodes() {
let child_handles = setup_nodes(13_000, 5, 5, false);
let child_handles = setup_nodes(5, 5, false);
assert_eq!(child_handles.len(), 5);

// Assert that we can set a key.
let first_non_seed_node = &child_handles[1];
let node_base_url = first_non_seed_node.base_url();
let set_kv_response =
set_kv("http://127.0.0.1:13001/set_kv", "some_key", "some_value").unwrap();
set_kv(&format!("{node_base_url}/set_kv"), "some_key", "some_value").unwrap();
assert_eq!(set_kv_response.status, true);

// Check node states through api.
let info = get_node_info("http://127.0.0.1:13001").unwrap();
let info = get_node_info(&node_base_url).unwrap();
assert_eq!(info.cluster_id, "testing");
assert_eq!(info.live_nodes.len(), 5);
assert_eq!(info.dead_nodes.len(), 0);
Expand All @@ -85,9 +136,11 @@ fn test_multiple_nodes() {

#[test]
fn test_multiple_nodes_with_dns_resolution_for_seed() {
let _child_handles = setup_nodes(12_000, 5, 5, true);
let child_handles = setup_nodes(5, 5, true);
let first_non_seed_node = &child_handles[1];

// Check node states through api.
let info = get_node_info("http://127.0.0.1:12001").unwrap();
let info = get_node_info(&first_non_seed_node.base_url()).unwrap();
assert!(info.cluster_state.node_state_snapshots.get(3).is_some());
assert_eq!(info.cluster_id, "testing");
assert_eq!(info.live_nodes.len(), 5);
Expand Down