use std::collections::HashMap; use std::pin::Pin; use super::behaviour::{CallTraceBehaviour, MockBehaviour}; use futures::stream::Stream; use futures::task::{Context, Poll}; use libp2p::swarm::handler::ConnectionHandler; use libp2p::swarm::{IntoConnectionHandler, NetworkBehaviour, Swarm, SwarmBuilder, SwarmEvent}; use libp2p::{PeerId, Transport}; use futures::StreamExt; pub fn new_test_swarm(behaviour: B) -> Swarm where B: NetworkBehaviour, { let id_keys = libp2p::identity::Keypair::generate_ed25519(); let local_public_key = id_keys.public(); let transport = libp2p::core::transport::MemoryTransport .upgrade(libp2p::core::upgrade::Version::V1) .authenticate(libp2p::plaintext::PlainText2Config { local_public_key: local_public_key.clone(), }) .multiplex(libp2p::yamux::YamuxConfig::default()) .boxed(); SwarmBuilder::new(transport, behaviour, local_public_key.into()).build() } pub fn random_multiaddr() -> libp2p::multiaddr::Multiaddr { libp2p::multiaddr::Protocol::Memory(rand::random::()).into() } /// Bind a memory multiaddr to a compatible swarm. pub async fn bind_listener( swarm: &mut Swarm, ) -> libp2p::multiaddr::Multiaddr { swarm.listen_on(random_multiaddr()).unwrap(); match swarm.select_next_some().await { SwarmEvent::NewListenAddr { listener_id: _, address, } => address, _ => panic!("Testing swarm's first event should be a new listener"), } } #[derive(Default)] pub struct SwarmPool { swarms: HashMap>, } impl SwarmPool { pub fn with_capacity(capacity: usize) -> Self { Self { swarms: HashMap::with_capacity(capacity), } } pub fn insert(&mut self, swarm: Swarm) -> PeerId { let peer_id = *swarm.local_peer_id(); self.swarms.insert(peer_id, swarm); peer_id } pub fn remove(&mut self, peer_id: &PeerId) { self.swarms.remove(peer_id); } pub fn get_mut(&mut self, peer_id: &PeerId) -> Option<&mut Swarm> { self.swarms.get_mut(peer_id) } pub fn swarms(&self) -> &HashMap> { &self.swarms } pub fn swarms_mut(&mut self) -> &mut HashMap> { &mut self.swarms } } impl Stream for SwarmPool where B: NetworkBehaviour, ::ConnectionHandler: ConnectionHandler, { type Item = (PeerId, SwarmEvent<::OutEvent, <<::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::Error>); fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut polls = self .get_mut() .swarms .iter_mut() .map(|(&peer_id, swarm)| swarm.map(move |ev| (peer_id, ev))) .collect::>(); polls.poll_next_unpin(cx) } }