Support expanding onto new subnet locally (#7086)

backend/canisters/registry/api/src/queries/subnets.rs

@@ -1,15 +1,15 @@
-use candid::Principal;
+use candid::{CandidType, Principal};
 use serde::{Deserialize, Serialize};
 use types::{CanisterId, Empty};
 
 pub type Args = Empty;
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
 pub enum Response {
     Success(Vec<Subnet>),
 }
 
-#[derive(Serialize, Deserialize, Clone, Debug)]
+#[derive(CandidType, Serialize, Deserialize, Clone, Debug)]
 pub struct Subnet {
     pub subnet_id: Principal,
     pub local_user_index: CanisterId,

backend/canisters/registry/api/src/updates/expand_onto_subnet.rs

@@ -1,10 +1,14 @@
 use candid::{CandidType, Principal};
 use human_readable::{HumanReadablePrincipal, ToHumanReadable};
 use serde::{Deserialize, Serialize};
+use types::CanisterId;
 
 #[derive(CandidType, Serialize, Deserialize, Debug)]
 pub struct Args {
     pub subnet_id: Principal,
+    pub local_user_index: Option<CanisterId>,
+    pub local_group_index: Option<CanisterId>,
+    pub notifications_canister: Option<CanisterId>,
 }
 
 #[derive(CandidType, Serialize, Deserialize, Debug)]
@@ -17,6 +21,12 @@ pub enum Response {
 #[derive(Serialize)]
 pub struct HumanReadableArgs {
     subnet_id: HumanReadablePrincipal,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    local_user_index: Option<HumanReadablePrincipal>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    local_group_index: Option<HumanReadablePrincipal>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    notifications_canister: Option<HumanReadablePrincipal>,
 }
 
 impl ToHumanReadable for Args {
@@ -25,6 +35,9 @@ impl ToHumanReadable for Args {
     fn to_human_readable(&self) -> Self::Target {
         HumanReadableArgs {
             subnet_id: self.subnet_id.into(),
+            local_user_index: self.local_user_index.map(HumanReadablePrincipal::from),
+            local_group_index: self.local_user_index.map(HumanReadablePrincipal::from),
+            notifications_canister: self.local_user_index.map(HumanReadablePrincipal::from),
         }
     }
 }

backend/canisters/registry/client/src/lib.rs

@@ -1,5 +1,9 @@
-use canister_client::generate_update_call;
+use canister_client::{generate_query_call, generate_update_call};
 use registry_canister::*;
 
+// Queries
+generate_query_call!(subnets);
+
 // Updates
 generate_update_call!(add_token);
+generate_update_call!(expand_onto_subnet);

backend/canisters/registry/impl/src/queries/subnets.rs

@@ -1,9 +1,9 @@
 use crate::{read_state, RuntimeState};
-use canister_api_macros::query;
 use canister_tracing_macros::trace;
+use ic_cdk::query;
 use registry_canister::subnets::{Response::*, *};
 
-#[query(msgpack = true)]
+#[query]
 #[trace]
 fn subnets(_args: Args) -> Response {
     read_state(subnets_impl)

backend/canisters/registry/impl/src/updates/expand_onto_subnet.rs

@@ -24,7 +24,19 @@ fn expand_onto_subnet_impl(args: Args, state: &mut RuntimeState) -> Result<Expan
     } else if state.data.subnets.in_progress().is_some() {
         Err(AlreadyInProgress)
     } else {
-        state.data.subnets.start_new(args.subnet_id, state.env.now());
+        let now = state.env.now();
+        state.data.subnets.start_new(args.subnet_id, now);
+
+        if state.data.test_mode {
+            state.data.subnets.update_in_progress(
+                |s| {
+                    s.local_user_index = args.local_user_index;
+                    s.local_group_index = args.local_group_index;
+                    s.notifications_canister = args.notifications_canister;
+                },
+                now,
+            );
+        }
 
         Ok(ExpandOntoSubnetJob {
             subnet_id: args.subnet_id,

backend/integration_tests/src/client/registry.rs

@@ -1,8 +1,8 @@
-use crate::{generate_msgpack_query_call, generate_msgpack_update_call, generate_update_call};
+use crate::{generate_msgpack_query_call, generate_msgpack_update_call, generate_query_call, generate_update_call};
 use registry_canister::*;
 
 // Queries
-generate_msgpack_query_call!(subnets);
+generate_query_call!(subnets);
 generate_msgpack_query_call!(updates);
 
 // Updates
@@ -26,7 +26,17 @@ pub mod happy_path {
         registry_canister_id: CanisterId,
         subnet_id: Principal,
     ) -> Subnet {
-        let response = super::expand_onto_subnet(env, sender, registry_canister_id, &expand_onto_subnet::Args { subnet_id });
+        let response = super::expand_onto_subnet(
+            env,
+            sender,
+            registry_canister_id,
+            &expand_onto_subnet::Args {
+                subnet_id,
+                local_user_index: None,
+                local_group_index: None,
+                notifications_canister: None,
+            },
+        );
 
         assert!(matches!(response, expand_onto_subnet::Response::Success));
 

backend/tools/canister_installer/Cargo.toml

@@ -28,6 +28,7 @@ notifications_index_canister_client = { path = "../../canisters/notifications_in
 online_users_canister = { path = "../../canisters/online_users/api" }
 proposals_bot_canister = { path = "../../canisters/proposals_bot/api" }
 registry_canister = { path = "../../canisters/registry/api" }
+registry_canister_client = { path = "../../canisters/registry/client" }
 sha256 = { path = "../../libraries/sha256" }
 sign_in_with_email_canister = { workspace = true }
 sign_in_with_email_canister_test_utils = { workspace = true }

backend/tools/canister_installer/src/lib.rs

@@ -50,18 +50,14 @@ async fn install_service_canisters_impl(
         set_controllers(
             management_canister,
             &canister_ids.local_user_index,
-            vec![canister_ids.user_index],
+            vec![canister_ids.registry],
         ),
         set_controllers(
             management_canister,
             &canister_ids.local_group_index,
-            vec![canister_ids.group_index],
-        ),
-        set_controllers(
-            management_canister,
-            &canister_ids.notifications,
-            vec![canister_ids.notifications_index],
+            vec![canister_ids.registry],
         ),
+        set_controllers(management_canister, &canister_ids.notifications, vec![canister_ids.registry]),
     ])
     .await;
 
@@ -552,59 +548,29 @@ async fn install_service_canisters_impl(
     .await
     .unwrap();
 
-    let add_local_group_index_canister_response = group_index_canister_client::add_local_group_index_canister(
-        agent,
-        &canister_ids.group_index,
-        &group_index_canister::add_local_group_index_canister::Args {
-            canister_id: canister_ids.local_group_index,
-            local_user_index_canister_id: canister_ids.local_user_index,
-            notifications_canister_id: canister_ids.notifications,
-        },
-    )
-    .await
-    .unwrap();
-
-    if !matches!(
-        add_local_group_index_canister_response,
-        group_index_canister::add_local_group_index_canister::Response::Success
-    ) {
-        panic!("{add_local_group_index_canister_response:?}");
-    }
-
-    let add_local_user_index_canister_response = user_index_canister_client::add_local_user_index_canister(
+    registry_canister_client::expand_onto_subnet(
         agent,
-        &canister_ids.user_index,
-        &user_index_canister::add_local_user_index_canister::Args {
-            canister_id: canister_ids.local_user_index,
-            notifications_canister_id: canister_ids.notifications,
+        &canister_ids.registry,
+        &registry_canister::expand_onto_subnet::Args {
+            subnet_id: Principal::anonymous(),
+            local_user_index: Some(canister_ids.local_user_index),
+            local_group_index: Some(canister_ids.local_group_index),
+            notifications_canister: Some(canister_ids.notifications),
         },
     )
     .await
     .unwrap();
 
-    if !matches!(
-        add_local_user_index_canister_response,
-        user_index_canister::add_local_user_index_canister::Response::Success
-    ) {
-        panic!("{add_local_user_index_canister_response:?}");
-    }
-
-    let add_notifications_canister_response = notifications_index_canister_client::add_notifications_canister(
-        agent,
-        &canister_ids.notifications_index,
-        &notifications_index_canister::add_notifications_canister::Args {
-            canister_id: canister_ids.notifications,
-            authorizers: vec![canister_ids.local_user_index, canister_ids.local_group_index],
-        },
-    )
-    .await
-    .unwrap();
+    for _ in 0..20 {
+        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+        let registry_canister::subnets::Response::Success(subnets) =
+            registry_canister_client::subnets(agent, &canister_ids.registry, &registry_canister::subnets::Args {})
+                .await
+                .unwrap();
 
-    if !matches!(
-        add_notifications_canister_response,
-        notifications_index_canister::add_notifications_canister::Response::Success
-    ) {
-        panic!("{add_notifications_canister_response:?}");
+        if !subnets.is_empty() {
+            break;
+        }
     }
 
     println!("Canister wasms installed");