feat(PocketIC): ability to create all mainnet subnets

packages/pocket-ic/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - The function `PocketIcBuilder::with_mainnet_nns_subnet_id` to specify that the NNS subnet should be created with the mainnet NNS subnet ID.
+- The function `PocketIcBuilder::with_all_mainnet_subnets` to specify that all mainnet subnets should be created.
 
 
 

packages/pocket-ic/src/common/rest.rs

@@ -537,6 +537,7 @@ impl From<SubnetConfigSet> for ExtendedSubnetConfigSet {
             system: vec![SubnetSpec::default(); system],
             application: vec![SubnetSpec::default(); application],
             verified_application: vec![SubnetSpec::default(); verified_application],
+            all_mainnet_subnets: None,
         }
     }
 }
@@ -659,6 +660,7 @@ pub struct ExtendedSubnetConfigSet {
     pub system: Vec<SubnetSpec>,
     pub application: Vec<SubnetSpec>,
     pub verified_application: Vec<SubnetSpec>,
+    pub all_mainnet_subnets: Option<bool>,
 }
 
 /// Specifies various configurations for a subnet.
@@ -770,6 +772,7 @@ impl ExtendedSubnetConfigSet {
             || self.ii.is_some()
             || self.fiduciary.is_some()
             || self.bitcoin.is_some()
+            || self.all_mainnet_subnets.unwrap_or_default()
         {
             return Ok(());
         }

packages/pocket-ic/src/lib.rs

@@ -492,6 +492,13 @@ impl PocketIcBuilder {
         self.mainnet_nns_subnet_id = Some(true);
         self
     }
+
+    pub fn with_all_mainnet_subnets(mut self) -> Self {
+        let mut config = self.config.unwrap_or_default();
+        config.all_mainnet_subnets = Some(true);
+        self.config = Some(config);
+        self
+    }
 }
 
 /// Representation of system time as duration since UNIX epoch

packages/pocket-ic/tests/tests.rs

@@ -412,7 +412,7 @@ fn test_initial_timestamp_with_cycles_minting() {
     // Initial time is bumped during each subnet creation and when executing rounds to deploy the CMC.
     assert_eq!(
         pic.get_time().as_nanos_since_unix_epoch(),
-        initial_timestamp + 7
+        initial_timestamp + 5
     );
 }
 
@@ -3273,3 +3273,42 @@ fn mainnet_nns_subnet_id() {
             .unwrap()
     );
 }
+
+#[test]
+fn all_mainnet_subnets() {
+    let pic = PocketIcBuilder::new().with_all_mainnet_subnets().build();
+
+    assert_eq!(pic.topology().subnet_configs.len(), 47);
+}
+
+#[test]
+fn nns_and_all_mainnet_subnets() {
+    let state = PocketIcState::new();
+    let icp_features = IcpFeatures {
+        cycles_minting: Some(IcpFeaturesConfig::DefaultConfig),
+        ..Default::default()
+    };
+    let pic = PocketIcBuilder::new()
+        .with_nns_subnet()
+        .with_icp_features(icp_features)
+        .with_state(state)
+        .build();
+    let topology = pic.topology();
+    let state = pic.drop_and_take_state().unwrap();
+
+    let nns_subnet = topology.get_nns().unwrap();
+    let nns_subnet_seed = topology
+        .subnet_configs
+        .get(&nns_subnet)
+        .unwrap()
+        .subnet_seed;
+    let state_dir = state.into_path();
+    let nns_state_dir = state_dir.join(hex::encode(nns_subnet_seed));
+
+    let pic = PocketIcBuilder::new()
+        .with_nns_state(nns_state_dir)
+        .with_all_mainnet_subnets()
+        .build();
+
+    assert_eq!(pic.topology().subnet_configs.len(), 47);
+}

rs/pocket_ic_server/CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - The endpoint `/instances/` takes an additional optional field `mainnet_nns_subnet_id` specifying that the NNS subnet should be created with the mainnet NNS subnet ID.
+- The type `ExtendedSubnetConfigSet` of the field `subnet_config_set` in the endpoint `/instances/` contains an additional optional field `all_mainnet_subnets` specifying that all mainnet subnets should be created.
 
 ### Changed
 - All subnets with mainnet canister ranges but the NNS subnet are always created with mainnet subnet IDs.

rs/pocket_ic_server/src/pocket_ic.rs

@@ -991,6 +991,19 @@ impl PocketIcSubnets {
             self.persist_registry_changes();
         }
 
+        let subnet_config = SubnetConfigInternal {
+            subnet_id,
+            subnet_kind,
+            instruction_config,
+            ranges,
+            alloc_range,
+        };
+        self.subnet_configs.push(subnet_config.clone());
+
+        Ok(subnet_config)
+    }
+
+    fn sync_subnet_time(&self) {
         // All subnets must have the same time and time can only advance =>
         // set the time to the maximum time in the latest state across all subnets.
         let mut time: SystemTime = self.initial_time;
@@ -1008,19 +1021,10 @@ impl PocketIcSubnets {
             subnet.state_machine.set_time(time);
             subnet.state_machine.execute_round();
         }
+    }
 
-        let subnet_config = SubnetConfigInternal {
-            subnet_id,
-            subnet_kind,
-            instruction_config,
-            ranges,
-            alloc_range,
-        };
-        self.subnet_configs.push(subnet_config.clone());
-
-        if let Some(icp_features) = self.icp_features.clone()
-            && update_registry_and_system_canisters
-        {
+    fn deploy_icp_features(&mut self, newly_created_subnets: Vec<SubnetId>) {
+        if let Some(icp_features) = self.icp_features.clone() {
             // using `let IcpFeatures { }` with explicit field names
             // to force an update after adding a new field to `IcpFeatures`
             let IcpFeatures {
@@ -1040,7 +1044,7 @@ impl PocketIcSubnets {
                 self.update_registry(config);
             }
             if let Some(ref config) = cycles_minting {
-                self.update_cmc(config, &subnet_kind);
+                self.update_cmc(config, newly_created_subnets);
             }
             if let Some(ref config) = icp_token {
                 self.deploy_icp_token(config);
@@ -1070,8 +1074,6 @@ impl PocketIcSubnets {
                 self.deploy_canister_migration(config);
             }
         }
-
-        Ok(subnet_config)
     }
 
     fn get_nns(&self) -> Option<Arc<StateMachine>> {
@@ -1162,7 +1164,7 @@ impl PocketIcSubnets {
         self.synced_registry_version = self.registry_data_provider.latest_version();
     }
 
-    fn update_cmc(&mut self, config: &IcpFeaturesConfig, subnet_kind: &SubnetKind) {
+    fn update_cmc(&mut self, config: &IcpFeaturesConfig, newly_created_subnets: Vec<SubnetId>) {
         // Using a match here to force an update after changing
         // the type of `IcpFeaturesConfig`.
         match config {
@@ -1303,13 +1305,14 @@ impl PocketIcSubnets {
         );
 
         // add fiduciary subnet to CMC
-        if let SubnetKind::Fiduciary = subnet_kind {
-            let fiduciary_subnet_id = self
-                .subnet_configs
-                .iter()
-                .find(|subnet_config| matches!(subnet_config.subnet_kind, SubnetKind::Fiduciary))
-                .map(|subnet_config| subnet_config.subnet_id)
-                .unwrap();
+        let fiduciary_subnet_id = self
+            .subnet_configs
+            .iter()
+            .find(|subnet_config| matches!(subnet_config.subnet_kind, SubnetKind::Fiduciary))
+            .map(|subnet_config| subnet_config.subnet_id);
+        if let Some(fiduciary_subnet_id) = fiduciary_subnet_id
+            && newly_created_subnets.contains(&fiduciary_subnet_id)
+        {
             let update_subnet_type_args = UpdateSubnetTypeArgs::Add("fiduciary".to_string());
             // returns ()
             self.execute_ingress_on(
@@ -2881,6 +2884,36 @@ impl PocketIc {
             subnet_config_info
         };
 
+        if subnet_configs.all_mainnet_subnets.unwrap_or_default() {
+            let mut subnet_ids_to_ranges = BTreeMap::new();
+            for (canister_range, subnet_id) in mainnet_routing_table.iter() {
+                subnet_ids_to_ranges
+                    .entry(subnet_id)
+                    .or_insert_with(Vec::new)
+                    .push(canister_range);
+            }
+            let existing_ranges: Vec<CanisterIdRange> = subnet_config_info
+                .iter()
+                .flat_map(|info| info.ranges.clone())
+                .collect();
+            for (subnet_id, ranges) in subnet_ids_to_ranges {
+                let already_exists = existing_ranges
+                    .iter()
+                    .any(|range| range.contains(&ranges[0].start));
+                if !already_exists {
+                    subnet_config_info.push(SubnetConfigInfo {
+                        ranges: ranges.into_iter().cloned().collect(),
+                        alloc_range: Some(range_gen.next_range()),
+                        subnet_id: Some(*subnet_id),
+                        subnet_state_dir: None,
+                        subnet_kind: SubnetKind::Application,
+                        instruction_config: SubnetInstructionConfig::Production,
+                        expected_state_time: None,
+                    });
+                }
+            }
+        }
+
         // NNS subnet must be sorted first
         subnet_config_info.sort_by(|subnet_config_info1, subnet_config_info2| {
             let non_nns1 = !matches!(subnet_config_info1.subnet_kind, SubnetKind::NNS);
@@ -2919,6 +2952,17 @@ impl PocketIc {
             subnet_configs.push(subnet_config_internal);
         }
 
+        subnets.sync_subnet_time();
+        if update_registry_and_system_canisters {
+            subnets.deploy_icp_features(
+                subnets
+                    .get_all()
+                    .into_iter()
+                    .map(|subnet| subnet.get_subnet_id())
+                    .collect(),
+            );
+        }
+
         let default_effective_canister_id = subnet_configs
             .iter()
             .find(|config| config.subnet_kind == SubnetKind::Application)
@@ -5186,6 +5230,8 @@ fn route(
                         },
                         update_registry_and_system_canisters,
                     )?;
+                    pic.subnets.sync_subnet_time();
+                    pic.subnets.deploy_icp_features(vec![subnet_id]);
                     pic.subnets
                         .persist_topology(pic.default_effective_canister_id);
                     Ok(pic.try_route_canister(canister_id).unwrap())

rs/sns/testing/README.md

@@ -61,7 +61,7 @@ by `sns-testing` identity will be automatically adopted.
 Additionally, `sns-testing-init` will output the ID of the NNS neuron that should be used to create NNS proposals.
 ```
 ...
-Use the following NNS neuron ID for further testing: 3912484856864073044
+Use the following NNS neuron ID for further testing: 350573795712518463
 ```
 
 This is the preconfigured NNS neuron controlled by identity specified via `--dev-identity`, and it is expected to have a lot of voting power. This neuron ID will be used later to submit NNS proposals that get adapted right away.
@@ -117,7 +117,7 @@ To run the basic scenario on the local IC network instance:
 2) Launch the basic SNS testing scenario:
    ```
    sns-testing --network http://127.0.0.1:8080 run-basic-scenario \
-      --dev-identity sns-testing --nns-neuron-id 3912484856864073044 \
+      --dev-identity sns-testing --nns-neuron-id 350573795712518463 \
       --canister-id "$(dfx canister --network http://127.0.0.1:8080 id test)" \
       --upgrade-wasm-path "$(bazel info bazel-bin)/rs/sns/testing/sns_testing_canister.wasm.gz" \
       --upgrade-candid-arg '(record { greeting = "Hi" })'
@@ -198,7 +198,7 @@ The example will use `//rs/sns/testing:sns_testing_canister` canister as SNS-con
    ```
    # sns doesn't support CLI-provided identities despite '--identity' option
    dfx identity use sns-testing
-   sns propose --network http://127.0.0.1:8080 --neuron-id 3912484856864073044 "$PWD/sns_init.yaml"
+   sns propose --network http://127.0.0.1:8080 --neuron-id 350573795712518463 "$PWD/sns_init.yaml"
    ```
 
 4) Complete the swap for the newly created SNS

rs/sns/testing/tests/sns_testing_ci.rs

@@ -154,7 +154,7 @@ async fn test_sns_testing_dev_nns_neuron_id() {
 
     // The following NNS neuron ID is hard-coded in sns-testing README.
     // If the test fails, then the README must be updated.
-    assert_eq!(dev_nns_neuron_id.id, 3912484856864073044);
+    assert_eq!(dev_nns_neuron_id.id, 350573795712518463);
 }
 
 #[tokio::test]