feat(rollout-controller): Implementing action taking and refactoring tests

Cargo.Bazel.lock

@@ -1,5 +1,5 @@
 {
-  "checksum": "683ef224e0a517f71f7b62868192a1d7a020d93d122f65348365ef7029c150ee",
+  "checksum": "3e505218ac8b9878195c82edc086ac1c9c01901eaf619e2cc955439273525d79",
   "crates": {
     "actix-codec 0.5.2": {
       "name": "actix-codec",
@@ -11519,6 +11519,15 @@
       "version": "0.3.0",
       "repository": null,
       "targets": [
+        {
+          "Library": {
+            "crate_name": "dre",
+            "crate_root": "src/lib.rs",
+            "srcs": [
+              "**/*.rs"
+            ]
+          }
+        },
         {
           "BuildScript": {
             "crate_name": "build_script_build",
@@ -11529,7 +11538,7 @@
           }
         }
       ],
-      "library_target_name": null,
+      "library_target_name": "dre",
       "common_attrs": {
         "compile_data_glob": [
           "**"

rs/cli/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@crate_index_dre//:defs.bzl", "aliases", "all_crate_deps")
-load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test")
+load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test", "rust_library")
 
 DEPS = [
     "//rs/ic-canisters",
@@ -8,6 +8,8 @@ DEPS = [
     "//rs/ic-management-backend:ic-management-backend-lib",
 ]
 
+package(default_visibility = ["//visibility:public"])
+
 rust_binary(
     name = "dre",
     srcs = glob(["src/**/*.rs"]),
@@ -16,6 +18,20 @@ rust_binary(
     proc_macro_deps = all_crate_deps(
         proc_macro = True,
     ),
+    deps = all_crate_deps(
+        normal = True,
+    ) + DEPS + ["//rs/cli:dre-lib"],
+)
+
+rust_library(
+    name = "dre-lib",
+    srcs = glob(["src/**/*.rs"]),
+    aliases = aliases(),
+    compile_data = glob(["config/**/*"]),
+    crate_name = "dre",
+    proc_macro_deps = all_crate_deps(
+        proc_macro = True,
+    ),
     deps = all_crate_deps(
         normal = True,
     ) + DEPS,

rs/cli/Cargo.toml

@@ -62,3 +62,11 @@ tempfile = "3.10.0"
 
 [dev-dependencies]
 wiremock = "0.6.0"
+
+[[bin]]
+name = "dre"
+path = "src/main.rs"
+
+[lib]
+name = "dre"
+path = "src/lib.rs"

rs/cli/src/cli.rs

@@ -9,45 +9,45 @@ use log::error;
 use crate::detect_neuron::{detect_hsm_auth, detect_neuron, Auth, Neuron};
 
 // For more info about the version setup, look at https://docs.rs/clap/latest/clap/struct.Command.html#method.version
-#[derive(Parser, Clone)]
+#[derive(Parser, Clone, Default)]
 #[clap(about, version = env!("CARGO_PKG_VERSION"), author)]
 pub struct Opts {
     #[clap(long, env = "HSM_PIN", global = true)]
-    pub(crate) hsm_pin: Option<String>,
+    pub hsm_pin: Option<String>,
     #[clap(long, value_parser=maybe_hex::<u64>, env = "HSM_SLOT", global = true)]
-    pub(crate) hsm_slot: Option<u64>,
+    pub hsm_slot: Option<u64>,
     #[clap(long, env = "HSM_KEY_ID", global = true)]
-    pub(crate) hsm_key_id: Option<String>,
+    pub hsm_key_id: Option<String>,
     #[clap(long, env = "PRIVATE_KEY_PEM", global = true)]
-    pub(crate) private_key_pem: Option<String>,
+    pub private_key_pem: Option<String>,
     #[clap(long, env = "NEURON_ID", global = true)]
-    pub(crate) neuron_id: Option<u64>,
+    pub neuron_id: Option<u64>,
     #[clap(long, env = "IC_ADMIN", global = true)]
-    pub(crate) ic_admin: Option<String>,
+    pub ic_admin: Option<String>,
     #[clap(long, env = "DEV", global = true)]
-    pub(crate) dev: bool,
+    pub dev: bool,
 
     // Skip the confirmation prompt
     #[clap(short, long, env = "YES", global = true, conflicts_with = "simulate")]
-    pub(crate) yes: bool,
+    pub yes: bool,
 
     // Simulate submission of the proposal, but do not actually submit it.
     #[clap(long, aliases = ["dry-run", "dryrun", "no"], global = true, conflicts_with = "yes")]
-    pub(crate) simulate: bool,
+    pub simulate: bool,
 
     #[clap(long, env = "VERBOSE", global = true)]
-    pub(crate) verbose: bool,
+    pub verbose: bool,
 
     // Specify the target network: "mainnet" (default), "staging", or NNS URL
     #[clap(long, env = "NETWORK", default_value = "mainnet")]
-    pub(crate) network: Network,
+    pub network: Network,
 
     #[clap(subcommand)]
-    pub(crate) subcommand: Commands,
+    pub subcommand: Commands,
 }
 
 #[derive(Subcommand, Clone)]
-pub(crate) enum Commands {
+pub enum Commands {
     // Convert a DER file to a Principal
     DerToPrincipal {
         /// Path to the DER file
@@ -155,7 +155,13 @@ pub(crate) enum Commands {
     },
 }
 
-pub(crate) mod subnet {
+impl Default for Commands {
+    fn default() -> Self {
+        Commands::Get { args: vec![] }
+    }
+}
+
+pub mod subnet {
     use super::*;
     use ic_base_types::PrincipalId;
 
@@ -269,7 +275,7 @@ pub(crate) mod subnet {
     }
 }
 
-pub(crate) mod version {
+pub mod version {
     use super::*;
 
     #[derive(Subcommand, Clone)]
@@ -322,7 +328,7 @@ pub(crate) mod version {
     }
 }
 
-pub(crate) mod hostos {
+pub mod hostos {
     use crate::operations::hostos_rollout::{NodeAssignment, NodeOwner};
 
     use super::*;
@@ -367,7 +373,7 @@ pub(crate) mod hostos {
     }
 }
 
-pub(crate) mod nodes {
+pub mod nodes {
     use super::*;
 
     #[derive(Parser, Clone)]

rs/cli/src/general.rs

@@ -17,7 +17,7 @@ use url::Url;
 
 use crate::detect_neuron::{Auth, Neuron};
 
-pub(crate) async fn vote_on_proposals(
+pub async fn vote_on_proposals(
     neuron: &Neuron,
     nns_url: &Url,
     accepted_proposers: &[u64],
@@ -89,7 +89,7 @@ pub(crate) async fn vote_on_proposals(
     Ok(())
 }
 
-pub(crate) async fn get_node_metrics_history(
+pub async fn get_node_metrics_history(
     wallet: CanisterId,
     subnets: Vec<PrincipalId>,
     start_at_nanos: u64,

rs/cli/src/ic_admin.rs

@@ -190,12 +190,7 @@ impl IcAdminWrapper {
         );
     }
 
-    pub(crate) fn propose_run(
-        &self,
-        cmd: ProposeCommand,
-        opts: ProposeOptions,
-        simulate: bool,
-    ) -> anyhow::Result<String> {
+    pub fn propose_run(&self, cmd: ProposeCommand, opts: ProposeOptions, simulate: bool) -> anyhow::Result<String> {
         let exec = |cli: &IcAdminWrapper, cmd: ProposeCommand, opts: ProposeOptions, add_dryrun_arg: bool| {
             if let Some(summary) = opts.clone().summary {
                 let summary_count = summary.chars().count();
@@ -311,7 +306,7 @@ impl IcAdminWrapper {
         }
     }
 
-    pub(crate) fn run(&self, command: &str, args: &[String], with_auth: bool) -> anyhow::Result<String> {
+    pub fn run(&self, command: &str, args: &[String], with_auth: bool) -> anyhow::Result<String> {
         let ic_admin_args = [&[command.to_string()], args].concat();
         self._run_ic_admin_with_args(&ic_admin_args, with_auth)
     }
@@ -346,7 +341,7 @@ impl IcAdminWrapper {
     }
 
     /// Run an `ic-admin get-*` command directly, and without an HSM
-    pub(crate) fn run_passthrough_get(&self, args: &[String]) -> anyhow::Result<()> {
+    pub fn run_passthrough_get(&self, args: &[String]) -> anyhow::Result<()> {
         if args.is_empty() {
             println!("List of available ic-admin 'get' sub-commands:\n");
             for subcmd in self.grep_subcommands(r"\s+get-(.+?)\s") {
@@ -378,7 +373,7 @@ impl IcAdminWrapper {
     }
 
     /// Run an `ic-admin propose-to-*` command directly
-    pub(crate) fn run_passthrough_propose(&self, args: &[String], simulate: bool) -> anyhow::Result<()> {
+    pub fn run_passthrough_propose(&self, args: &[String], simulate: bool) -> anyhow::Result<()> {
         if args.is_empty() {
             println!("List of available ic-admin 'propose' sub-commands:\n");
             for subcmd in self.grep_subcommands(r"\s+propose-to-(.+?)\s") {
@@ -572,7 +567,7 @@ impl IcAdminWrapper {
         Ok((update_urls, expected_hash))
     }
 
-    pub(crate) async fn prepare_to_propose_to_update_elected_versions(
+    pub async fn prepare_to_propose_to_update_elected_versions(
         release_artifact: &Artifact,
         version: &String,
         release_tag: &String,
@@ -910,7 +905,7 @@ must be identical, and must match the SHA256 from the payload of the NNS proposa
 
 #[derive(Display, Clone)]
 #[strum(serialize_all = "kebab-case")]
-pub(crate) enum ProposeCommand {
+pub enum ProposeCommand {
     ChangeSubnetMembership {
         subnet_id: PrincipalId,
         node_ids_add: Vec<PrincipalId>,

rs/cli/src/lib.rs

@@ -0,0 +1,25 @@
+pub mod cli;
+pub mod clients;
+pub(crate) mod defaults;
+pub mod detect_neuron;
+pub mod general;
+pub mod ic_admin;
+pub mod operations;
+pub mod ops_subnet_node_replace;
+pub mod registry_dump;
+pub mod runner;
+
+/// Get a localhost socket address with random, unused port.
+pub fn local_unused_port() -> u16 {
+    let addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap();
+    let socket = socket2::Socket::new(
+        socket2::Domain::IPV4,
+        socket2::Type::STREAM,
+        Some(socket2::Protocol::TCP),
+    )
+    .unwrap();
+    socket.bind(&addr.into()).unwrap();
+    socket.set_reuse_address(true).unwrap();
+    let tcp = std::net::TcpListener::from(socket);
+    tcp.local_addr().unwrap().port()
+}

rs/cli/src/main.rs

@@ -1,8 +1,9 @@
-use crate::general::{get_node_metrics_history, vote_on_proposals};
 use crate::ic_admin::IcAdminWrapper;
-use crate::operations::hostos_rollout::{NodeGroupUpdate, NumberOfNodes};
 use clap::{error::ErrorKind, CommandFactory, Parser};
 use dotenv::dotenv;
+use dre::general::{get_node_metrics_history, vote_on_proposals};
+use dre::operations::hostos_rollout::{NodeGroupUpdate, NumberOfNodes};
+use dre::{cli, ic_admin, local_unused_port, registry_dump, runner};
 use ic_base_types::CanisterId;
 use ic_canisters::governance::governance_canister_version;
 use ic_management_backend::endpoints;
@@ -14,17 +15,6 @@ use std::str::FromStr;
 use std::sync::mpsc;
 use std::thread;
 
-mod cli;
-mod clients;
-pub(crate) mod defaults;
-mod detect_neuron;
-mod general;
-mod ic_admin;
-mod operations;
-mod ops_subnet_node_replace;
-mod registry_dump;
-mod runner;
-
 const STAGING_NEURON_ID: u64 = 49;
 
 #[tokio::main]
@@ -398,21 +388,6 @@ fn parse_min_nakamoto_coefficients(
     })
 }
 
-/// Get a localhost socket address with random, unused port.
-fn local_unused_port() -> u16 {
-    let addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap();
-    let socket = socket2::Socket::new(
-        socket2::Domain::IPV4,
-        socket2::Type::STREAM,
-        Some(socket2::Protocol::TCP),
-    )
-    .unwrap();
-    socket.bind(&addr.into()).unwrap();
-    socket.set_reuse_address(true).unwrap();
-    let tcp = std::net::TcpListener::from(socket);
-    tcp.local_addr().unwrap().port()
-}
-
 fn init_logger() {
     match std::env::var("RUST_LOG") {
         Ok(val) => std::env::set_var("LOG_LEVEL", val),

rs/cli/src/runner.rs

@@ -210,7 +210,7 @@ impl Runner {
         })
     }
 
-    pub(crate) async fn prepare_versions_to_retire(
+    pub async fn prepare_versions_to_retire(
         &self,
         release_artifact: &Artifact,
         edit_summary: bool,
@@ -282,7 +282,7 @@ impl Runner {
             })
     }
 
-    pub(crate) async fn hostos_rollout_nodes(
+    pub async fn hostos_rollout_nodes(
         &self,
         node_group: NodeGroupUpdate,
         version: &String,

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

@@ -569,6 +569,12 @@ pub enum Network {
     Url(url::Url),
 }
 
+impl Default for Network {
+    fn default() -> Self {
+        Network::Mainnet
+    }
+}
+
 impl Debug for Network {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self)

rs/rollout-controller/BUILD.bazel

@@ -6,7 +6,8 @@ load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test")
 DEPS = [
     "//rs/ic-observability/service-discovery",
     "//rs/ic-management-types",
-    "//rs/ic-management-backend:ic-management-backend-lib",]
+    "//rs/ic-management-backend:ic-management-backend-lib",
+    "//rs/cli:dre-lib"]
 
 
 rust_binary(

rs/rollout-controller/Cargo.toml

@@ -38,6 +38,7 @@ candid = { workspace = true }
 ic-base-types = { workspace = true }
 pretty_assertions = { workspace = true }
 itertools = { workspace = true }
+dre = { path = "../cli"}
 
 [dev-dependencies]
 rstest = "0.18.2"