Add tests and fix some encountered issues

Cargo.toml

@@ -175,6 +175,7 @@ tokio = { version = "1.2.0", features = ["full"] }
 url = "2.5.0"
 urlencoding = "2.1.0"
 warp = "0.3"
+wiremock = "0.6.0"
 
 
 [profile.release]

rs/cli/Cargo.toml

@@ -60,7 +60,7 @@ url = { workspace = true }
 tempfile = "3.10.0"
 
 [dev-dependencies]
-wiremock = "0.6.0"
+wiremock = { workspace = true }
 
 [[bin]]
 name = "dre"

rs/ic-management-backend/src/main.rs

@@ -14,13 +14,8 @@ use clap::Parser;
 use dotenv::dotenv;
 use url::Url;
 
-#[derive(Parser)]
-struct Args {}
-
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
-    let _args = Args::parse();
-
     dotenv().ok();
     std::env::set_var("RUST_LOG", "info");
     env_logger::init();

rs/ic-management-backend/src/prometheus.rs

@@ -2,8 +2,5 @@ use ic_management_types::Network;
 use prometheus_http_query::Client;
 
 pub fn client(network: &Network) -> Client {
-    match network.name.as_str() {
-        "mainnet" => Client::try_from("https://victoria.mainnet.dfinity.network/select/0/prometheus/").unwrap(),
-        _ => Client::try_from("https://victoria.testnet.dfinity.network/select/0/prometheus").unwrap(),
-    }
+    Client::try_from(network.get_prometheus_endpoint().as_str()).unwrap()
 }

rs/ic-management-types/Cargo.toml

@@ -25,5 +25,8 @@ url = { workspace = true }
 anyhow = { workspace = true }
 candid = { workspace = true }
 
+[dev-dependencies]
+wiremock = { workspace = true }
+
 [lib]
 path = "src/lib.rs"

rs/ic-management-types/src/lib.rs

@@ -592,7 +592,7 @@ impl Network {
                 },
             ),
             "staging" => (
-                "mainnet".to_string(),
+                "staging".to_string(),
                 if nns_urls.is_empty() {
                     vec![Url::from_str("http://[2600:3000:6100:200:5000:b0ff:fe8e:6b7b]:8080").unwrap()]
                 } else {
@@ -608,10 +608,11 @@ impl Network {
                 },
             ),
         };
-        Ok(Network {
-            name,
-            nns_urls: find_reachable_nns_urls(nns_urls).await,
-        })
+        let nns_urls = find_reachable_nns_urls(nns_urls).await;
+        if nns_urls.is_empty() {
+            return Err("No reachable NNS URLs provided".to_string());
+        }
+        Ok(Network { name, nns_urls })
     }
 
     pub fn get_nns_urls(&self) -> &Vec<Url> {
@@ -623,13 +624,13 @@ impl Network {
             .iter()
             .map(|url| url.to_string())
             .collect::<Vec<String>>()
-            .join(", ")
+            .join(",")
     }
 
     pub fn get_prometheus_endpoint(&self) -> Url {
         match self.name.as_str() {
             "mainnet" => "https://victoria.mainnet.dfinity.network/select/0/prometheus/",
-            _ => "https://victoria.testnet.dfinity.network/select/0/prometheus",
+            _ => "https://victoria.testnet.dfinity.network/select/0/prometheus/",
         }
         .parse()
         .expect("Couldn't parse url")
@@ -729,3 +730,129 @@ async fn find_reachable_nns_urls(nns_urls: Vec<Url>) -> Vec<Url> {
 
     Vec::new()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use wiremock::MockServer;
+
+    #[tokio::test]
+    async fn test_network_new_mainnet() {
+        let name = "mainnet";
+        let nns_urls = vec![];
+        let network = Network::new(name, &nns_urls).await.unwrap();
+
+        assert_eq!(network.name, "mainnet");
+        assert_eq!(network.get_nns_urls(), &vec![Url::from_str("https://ic0.app").unwrap()]);
+    }
+
+    #[tokio::test]
+    async fn test_network_new_mainnet_custom_url() {
+        let mock_server = MockServer::start().await;
+        let mock_server_url: Url = mock_server.uri().parse().unwrap();
+        let network = Network::new("mainnet", &vec![mock_server_url.clone()]).await.unwrap();
+
+        assert_eq!(network.name, "mainnet");
+        assert_eq!(network.get_nns_urls(), &vec![mock_server_url]);
+    }
+
+    #[tokio::test]
+    async fn test_network_new_mainnet_custom_and_invalid_url() {
+        let mock_server = MockServer::start().await;
+        let mock_server_url: Url = mock_server.uri().parse().unwrap();
+        let invalid_url1 = Url::from_str("https://unreachable.url1").unwrap();
+        let invalid_url2 = Url::from_str("https://unreachable.url2").unwrap();
+
+        let expected_nns_urls = vec![mock_server_url.clone()];
+
+        // Test with the invalid URL last
+        let network = Network::new("mainnet", &vec![mock_server_url.clone(), invalid_url1.clone()])
+            .await
+            .unwrap();
+
+        assert_eq!(network.name, "mainnet");
+        assert_eq!(network.get_nns_urls(), &expected_nns_urls);
+
+        // Test with the invalid URL first
+        let network = Network::new("mainnet", &vec![invalid_url1.clone(), mock_server_url.clone()])
+            .await
+            .unwrap();
+
+        assert_eq!(network.name, "mainnet");
+        assert_eq!(network.get_nns_urls(), &expected_nns_urls);
+
+        // Test with the valid URL in the middle
+        let network = Network::new("mainnet", &vec![invalid_url1, mock_server_url.clone(), invalid_url2])
+            .await
+            .unwrap();
+
+        assert_eq!(network.name, "mainnet");
+        assert_eq!(network.get_nns_urls(), &expected_nns_urls);
+    }
+
+    #[tokio::test]
+    async fn test_network_new_staging() {
+        let network = Network::new("staging", &vec![]).await.unwrap();
+
+        assert_eq!(network.name, "staging");
+        assert_eq!(
+            network.get_nns_urls(),
+            &vec![Url::from_str("http://[2600:3000:6100:200:5000:b0ff:fe8e:6b7b]:8080").unwrap()]
+        );
+    }
+
+    #[tokio::test]
+    async fn test_network_new_all_unreachable() {
+        let name = "custom";
+        let nns_urls = vec![Url::from_str("https://unreachable.url").unwrap()];
+        let network = Network::new(name, &nns_urls).await;
+
+        assert_eq!(network, Err("No reachable NNS URLs provided".to_string()));
+    }
+
+    #[test]
+    fn test_network_get_nns_urls_string() {
+        let nns_urls = vec![
+            Url::from_str("https://ic0.app").unwrap(),
+            Url::from_str("https://custom.nns").unwrap(),
+        ];
+        let network = Network {
+            name: "mainnet".to_string(),
+            nns_urls,
+        };
+
+        assert_eq!(network.get_nns_urls_string(), "https://ic0.app/,https://custom.nns/");
+    }
+
+    #[test]
+    fn test_network_get_prometheus_endpoint() {
+        let network = Network {
+            name: "mainnet".to_string(),
+            nns_urls: vec![],
+        };
+
+        assert_eq!(
+            network.get_prometheus_endpoint(),
+            Url::parse("https://victoria.mainnet.dfinity.network/select/0/prometheus/").unwrap()
+        );
+
+        let network = Network {
+            name: "some_testnet".to_string(),
+            nns_urls: vec![],
+        };
+        assert_eq!(
+            network.get_prometheus_endpoint(),
+            Url::parse("https://victoria.testnet.dfinity.network/select/0/prometheus/").unwrap()
+        );
+    }
+
+    #[test]
+    fn test_network_legacy_name() {
+        let network = Network {
+            name: "mainnet".to_string(),
+            nns_urls: vec![],
+        };
+
+        assert_eq!(network.legacy_name(), "mercury");
+    }
+}

rs/np-notifications/src/main.rs

@@ -46,7 +46,6 @@ use crate::registry::{start_registry_updater_loop, RegistryLoopConfig};
 use crate::router::Router;
 use crate::service_health::ServiceHealth;
 use crate::sink::{LogSink, Sink};
-use clap::Parser;
 
 mod health_check;
 mod nodes_status;

rs/slack-notifications/Cargo.toml

@@ -15,6 +15,7 @@ dotenv = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }
 ic-agent = { workspace = true }
+ic-management-types = { workspace = true }
 ic-nns-common = { workspace = true }
 ic-nns-constants = { workspace = true }
 ic-nns-governance = { workspace = true }

rs/slack-notifications/src/main.rs

@@ -1,3 +1,4 @@
+use ic_management_types::Network;
 use ic_nns_governance::pb::v1::{ListProposalInfo, ListProposalInfoResponse, ProposalInfo, ProposalStatus};
 
 use anyhow::Result;
@@ -12,16 +13,14 @@ use std::time::SystemTime;
 use tokio::time::{sleep, Duration};
 mod slack;
 use clap::Parser;
+use reqwest::Url;
 
 #[macro_use]
 extern crate lazy_static;
 
 #[derive(Deserialize)]
 struct Config {}
 
-#[derive(Parser)]
-struct Args {}
-
 // Time to wait for a new proposal after the last one was created before sending
 // out the Slack notification.
 const COOLING_PERIOD_SECS: u64 = 60;
@@ -30,17 +29,34 @@ const SLACK_URL_ENV: &str = "SLACK_URL";
 
 #[tokio::main]
 async fn main() {
-    let _args = Args::parse();
     std::env::set_var("RUST_LOG", "info");
     env_logger::init();
     dotenv::dotenv().ok();
 
-    let failed_proposals_handle = tokio::spawn(notify_for_failed_proposals());
-    let new_proposals_handle = tokio::spawn(notify_for_new_proposals());
+    let args = Cli::parse();
+    let target_network = ic_management_types::Network::new(args.network.clone(), &args.nns_urls)
+        .await
+        .expect("Failed to create network");
+
+    let failed_proposals_handle = tokio::spawn(notify_for_failed_proposals(target_network.clone()));
+    let new_proposals_handle = tokio::spawn(notify_for_new_proposals(target_network));
 
     futures::future::join_all(vec![failed_proposals_handle, new_proposals_handle]).await;
 }
 
+#[derive(Parser, Debug)]
+#[clap(about, version)]
+struct Cli {
+    // Target network. Can be one of: "mainnet", "staging", or an arbitrary "<testnet>" name
+    #[clap(long, env = "NETWORK", default_value = "mainnet")]
+    network: String,
+
+    // NNS_URLs for the target network, comma separated.
+    // The argument is mandatory for testnets, and is optional for mainnet and staging
+    #[clap(long, env = "NNS_URLS", aliases = &["registry-url", "nns-url"], value_delimiter = ',')]
+    pub nns_urls: Vec<Url>,
+}
+
 #[derive(Default)]
 pub struct ProposalCheckpointStore {
     file_path: String,
@@ -96,11 +112,10 @@ struct ProposalPoller {
 }
 
 impl ProposalPoller {
-    fn new() -> Self {
+    fn new(target_network: Network) -> Self {
+        let nns_url = target_network.get_nns_urls()[0].clone();
         let agent = Agent::builder()
-            .with_transport(
-                ReqwestHttpReplicaV2Transport::create("https://ic0.app").expect("failed to create transport"),
-            )
+            .with_transport(ReqwestHttpReplicaV2Transport::create(nns_url).expect("failed to create transport"))
             .build()
             .expect("failed to build the agent");
         Self { agent }
@@ -146,10 +161,10 @@ impl ProposalPoller {
     }
 }
 
-async fn notify_for_new_proposals() {
+async fn notify_for_new_proposals(target_network: Network) {
     let mut last_notified_proposal =
         ProposalCheckpointStore::new("new").expect("failed to initialize last notified proposal tracking");
-    let proposal_poller = ProposalPoller::new();
+    let proposal_poller = ProposalPoller::new(target_network);
     loop {
         info!("sleeping");
         sleep(Duration::from_secs(10)).await;
@@ -224,10 +239,10 @@ async fn notify_for_new_proposals() {
     }
 }
 
-async fn notify_for_failed_proposals() {
+async fn notify_for_failed_proposals(target_network: Network) {
     let mut checkpoint =
         ProposalCheckpointStore::new("failed").expect("failed to initialize last notified proposal tracking");
-    let proposal_poller = ProposalPoller::new();
+    let proposal_poller = ProposalPoller::new(target_network);
     loop {
         info!("checking for failed proposals");
         if let Ok(mut proposals) = proposal_poller.poll_not_executed_once().await {