n0_mainline/dht/
testnet.rs

1use super::Dht;
2
3/// Create a testnet of Dht nodes to run tests against instead of the real mainline network.
4#[derive(Debug)]
5pub struct Testnet {
6    /// bootstrapping nodes for this testnet.
7    pub bootstrap: Vec<String>,
8    /// all nodes in this testnet
9    pub nodes: Vec<Dht>,
10}
11
12#[allow(dead_code)]
13impl Testnet {
14    /// Create a new testnet with a certain size.
15    ///
16    /// Note: this network will be shutdown as soon as this struct
17    /// gets dropped, if you want the network to be `'static`, then
18    /// you should call [Self::leak].
19    ///
20    /// This will await until all nodes are [bootstrapped][Dht::bootstrapped].
21    pub async fn new(count: usize) -> Result<Testnet, std::io::Error> {
22        let testnet = Testnet::new_inner(count, false, None).await?;
23
24        for node in &testnet.nodes {
25            node.bootstrapped().await?;
26        }
27
28        Ok(testnet)
29    }
30
31    #[cfg(test)]
32    pub(crate) async fn new_without_signed_peers(count: usize) -> Result<Testnet, std::io::Error> {
33        let testnet = Testnet::new_inner(count, true, None).await?;
34
35        for node in &testnet.nodes {
36            node.bootstrapped().await?;
37        }
38
39        Ok(testnet)
40    }
41
42    #[cfg(test)]
43    pub(crate) async fn new_with_bootstrap(
44        count: usize,
45        bootstrap: &[String],
46    ) -> Result<Testnet, std::io::Error> {
47        let testnet = Testnet::new_inner(count, false, Some(bootstrap.to_vec())).await?;
48
49        for node in &testnet.nodes {
50            node.bootstrapped().await?;
51        }
52
53        Ok(testnet)
54    }
55
56    async fn new_inner(
57        count: usize,
58        disable_signed_peers: bool,
59        bootstrap: Option<Vec<String>>,
60    ) -> Result<Testnet, std::io::Error> {
61        let mut nodes: Vec<Dht> = vec![];
62        let mut bootstrap = bootstrap.unwrap_or_default();
63
64        for i in 0..count {
65            let mut builder = Dht::builder();
66
67            if disable_signed_peers {
68                #[cfg(test)]
69                builder.disable_signed_peers();
70            }
71
72            let node = builder.server_mode().bootstrap(&bootstrap).build()?;
73
74            if i == 0 {
75                let info = node.info().await?;
76                let addr = info.local_addr();
77
78                bootstrap.push(format!("127.0.0.1:{}", addr.port()));
79            }
80
81            nodes.push(node);
82        }
83
84        let testnet = Self { bootstrap, nodes };
85
86        Ok(testnet)
87    }
88
89    /// By default as soon as this testnet gets dropped,
90    /// all the nodes get dropped and the entire network is shutdown.
91    ///
92    /// This method uses [Box::leak] to keep nodes running, which is
93    /// useful if you need to keep running the testnet in the process
94    /// even if this struct gets dropped.
95    pub fn leak(&self) {
96        for node in self.nodes.clone() {
97            Box::leak(Box::new(node));
98        }
99    }
100}
n0_mainline/dht/testnet.rs

n0_mainline/dht/
testnet.rs
