Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
10.0 kB · 299 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300use std::{ops::Deref, sync::Arc, time::Duration};
use async_trait::async_trait;use n0_future::{BufferedStreamExt, StreamExt, stream};use testresult::TestResult;
use super::*;
struct TestNode { node: Arc<Node>, public: Arc<crate::public_key::PublicKey>, ticket: String, #[allow(dead_code)] private: Vec<u8>,}
impl Deref for TestNode { type Target = Node;
fn deref(&self) -> &Self::Target { &self.node }}
impl TestNode { /// Helper to create a test node with given config and handler mode. async fn new(handler: HandlerMode) -> TestResult<TestNode> { let config = Config { key: vec![0_u8; 32], // will be replaced topic: vec![0_u8; 32], // all nodes use the same topic max_send_duration: Duration::from_secs(10), disable_relay: false, }; Self::new_with_config(handler, config).await }
async fn new_with_config(handler: HandlerMode, mut config: Config) -> TestResult<TestNode> { let key = iroh::SecretKey::generate(&mut rand::rng()); let key = key.to_bytes().to_vec(); config.key = key.clone(); let node = Node::new_in_runtime(config, handler).await?; let public = node.node_id().await?; let ticket = node.ticket().await?; Ok(TestNode { node, private: key, public, ticket, }) }}
/// Helper to create multiple test nodes with given handler mode.async fn test_nodes( n: usize, handler: impl Fn(usize) -> HandlerMode + Send + Sync, disable_relay: bool,) -> TestResult<Vec<TestNode>> { const PAR: usize = 32; let modes = (0..n).map(handler).collect::<Vec<_>>();
let config = Config { key: vec![0_u8; 32], // will be replaced topic: vec![0_u8; 32], max_send_duration: Duration::from_secs(10), disable_relay, }; // create all nodes in parallel let nodes = stream::iter(modes) .map(|mode| TestNode::new_with_config(mode, config.clone())) .buffered_unordered(PAR) .collect::<Vec<_>>() .await; let nodes = nodes.into_iter().collect::<TestResult<Vec<_>>>()?; // join everyone to everyone let tickets = nodes.iter().map(|n| n.ticket.clone()).collect::<Vec<_>>(); let res = stream::iter(&nodes) .map(|n| n.join_peers(tickets.clone())) .buffered_unordered(PAR) .collect::<Vec<_>>() .await; res.into_iter().collect::<Result<Vec<_>, _>>()?; Ok(nodes)}
#[tokio::test]async fn one_node() -> TestResult<()> { tracing_subscriber::fmt::try_init().ok(); let node = TestNode::new(HandlerMode::Sender).await?.node; let write = node.node_scope(); let db = node.db(); println!("Ticket: {}", node.ticket().await?); write .put(Some(b"stream1".to_vec()), b"s".to_vec(), b"y".to_vec()) .await?; write .put(Some(b"stream2".to_vec()), b"s".to_vec(), b"y".to_vec()) .await?; let res = db.subscribe_with_opts(SubscribeOpts { filter: Filter::new(), mode: SubscribeMode::Both, }); while let Some(item) = res.next_raw().await? { if let SubscribeItem::CurrentDone = item { break; } println!("Got item: {item:?}"); } let res = db .iter_with_opts( Filter::new() .stream(b"stream1".to_vec()) .scope(node.node_id().await?), ) .await?; println!("Iter result: {res:?}"); Ok(())}
struct TestHandler<T> { info: T, sender: tokio::sync::mpsc::Sender<(T, String, Vec<u8>)>,}
impl<T> TestHandler<T> { fn new(info: T, sender: tokio::sync::mpsc::Sender<(T, String, Vec<u8>)>) -> Self { Self { info, sender } }}
#[async_trait]impl<T: Clone + Send + Sync + 'static> DataHandler for TestHandler<T> { async fn handle_data(&self, _from: Arc<public_key::PublicKey>, topic: String, data: Vec<u8>) { self.sender .send((self.info.clone(), topic, data)) .await .ok(); }}
#[tokio::test]async fn two_nodes_send_receive() -> TestResult<()> { tracing_subscriber::fmt::try_init().ok(); let (tx, mut rx) = tokio::sync::mpsc::channel(32); let handler = Arc::new(TestHandler::new((), tx)); let sender = TestNode::new(HandlerMode::Sender).await?; let receiver = TestNode::new(HandlerMode::Receiver(handler)).await?; // join the sender to the receiver. This will also configure the receiver endpoint to be able to dial the sender. receiver.join_peers(vec![sender.ticket.clone()]).await?; let stream = "teststream".to_string(); receiver .subscribe(stream.clone(), sender.public.clone()) .await?; sender.send_segment(stream, b"segment1".to_vec()).await?; let (_, stream, data) = rx.recv().await.expect("should get data"); assert_eq!(stream, "teststream"); assert_eq!(data, b"segment1".to_vec()); Ok(())}
#[tokio::test]async fn three_nodes_send_forward_receive() -> TestResult<()> { tracing_subscriber::fmt::try_init().ok(); let (tx, mut rx) = tokio::sync::mpsc::channel(32); let handler = Arc::new(TestHandler::new((), tx)); let sender = TestNode::new(HandlerMode::Sender).await?; let forwarder = TestNode::new(HandlerMode::Forwarder).await?; let receiver = TestNode::new(HandlerMode::Receiver(handler)).await?; // join everyone to everyone, so the receiver can reach the sender via the forwarder. let tickets = vec![ sender.ticket.clone(), forwarder.ticket.clone(), receiver.ticket.clone(), ]; receiver.join_peers(tickets.clone()).await?; forwarder.join_peers(tickets.clone()).await?; sender.join_peers(tickets).await?; let stream = "teststream".to_string(); receiver .subscribe(stream.clone(), forwarder.public.clone()) .await?; forwarder .subscribe(stream.clone(), sender.public.clone()) .await?; sender.send_segment(stream, b"segment1".to_vec()).await?; let (_, stream, data) = rx.recv().await.expect("should get data"); assert_eq!(stream, "teststream"); assert_eq!(data, b"segment1".to_vec()); Ok(())}
#[tokio::test]async fn meta_three_nodes_send_forward_receive() -> TestResult<()> { tracing_subscriber::fmt::try_init().ok(); let (tx, mut rx) = tokio::sync::mpsc::channel(32); let handler = Arc::new(TestHandler::new((), tx)); let sender = TestNode::new(HandlerMode::Sender).await?; let forwarder = TestNode::new(HandlerMode::Forwarder).await?; let receiver = TestNode::new(HandlerMode::Receiver(handler)).await?; // join everyone to everyone, so the receiver can reach the sender via the forwarder. let tickets = vec![ sender.ticket.clone(), forwarder.ticket.clone(), receiver.ticket.clone(), ]; receiver.join_peers(tickets.clone()).await?; forwarder.join_peers(tickets.clone()).await?; sender.join_peers(tickets).await?; let stream = "teststream".to_string(); receiver .subscribe(stream.clone(), forwarder.public.clone()) .await?; forwarder .subscribe(stream.clone(), sender.public.clone()) .await?; sender.send_segment(stream, b"segment1".to_vec()).await?; let (_, stream, data) = rx.recv().await.expect("should get data"); assert_eq!(stream, "teststream"); assert_eq!(data, b"segment1".to_vec()); let stream = receiver.db().subscribe(Filter::new()); while let Some(item) = stream.next_raw().await? { println!("{}", subscribe_item_debug(&item)); } Ok(())}
async fn broadcast( nsenders: usize, nforwarders: usize, nreceivers: usize, nmsgs: usize,) -> TestResult<()> { let (tx, mut rx) = tokio::sync::mpsc::channel(32); let ntotal = nsenders + nforwarders + nreceivers; let senders = 0..nsenders; let forwarders = nsenders..(nsenders + nforwarders); let receivers = (nsenders + nforwarders)..ntotal; let make_handler = |i: usize| { if senders.contains(&i) { HandlerMode::Sender } else if forwarders.contains(&i) { HandlerMode::Forwarder } else { HandlerMode::Receiver(Arc::new(TestHandler::new(i, tx.clone()))) } }; let nodes = test_nodes(ntotal, make_handler, true).await?; let senders = &nodes[senders]; let forwarders = &nodes[forwarders]; let receivers = &nodes[receivers]; let stream = "teststream".to_string(); // subscribe all forwarders to a sender, round robin for (i, forwarder) in forwarders.iter().enumerate() { let sender = &senders[i % senders.len()]; forwarder .subscribe(stream.clone(), sender.public.clone()) .await?; } // subscribe all receivers to a forwarder, round robin for (i, receiver) in receivers.iter().enumerate() { let forwarder = &forwarders[i % forwarders.len()]; receiver .subscribe(stream.clone(), forwarder.public.clone()) .await?; } for _ in 0..nmsgs { for sender in senders { sender .send_segment(stream.clone(), b"segment1".to_vec()) .await?; } for _ in 0..receivers.len() { let (i, stream, _) = rx.recv().await.expect("should get data"); println!("Node {i} got data on stream {stream}"); } } Ok(())}
#[tokio::test]async fn broadcast_1_2_4() -> TestResult<()> { tracing_subscriber::fmt().try_init().ok(); broadcast(1, 2, 4, 1).await?; Ok(())}
#[tokio::test]async fn broadcast_1_3_9() -> TestResult<()> { tracing_subscriber::fmt().try_init().ok(); broadcast(1, 3, 9, 1).await?; Ok(())}
#[tokio::test]async fn broadcast_1_4_16() -> TestResult<()> { tracing_subscriber::fmt().try_init().ok(); broadcast(1, 4, 16, 100).await?; Ok(())}