Add tests

packages/pocket-ic/src/lib.rs

@@ -1,4 +1,5 @@
 #![allow(clippy::test_attr_in_doctest)]
+use crate::common::rest::TickConfigs;
 /// # PocketIC: A Canister Testing Platform
 ///
 /// PocketIC is the local canister smart contract testing platform for the [Internet Computer](https://internetcomputer.org/).
@@ -522,6 +523,14 @@ impl PocketIc {
         runtime.block_on(async { self.pocket_ic.tick().await })
     }
 
+    /// Make the IC produce and progress by one block with custom
+    /// configs for the round.
+    #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id))]
+    pub fn tick_with_configs(&self, configs: TickConfigs) {
+        let runtime = self.runtime.clone();
+        runtime.block_on(async { self.pocket_ic.tick_with_configs(configs).await })
+    }
+
     /// Configures the IC to make progress automatically,
     /// i.e., periodically update the time of the IC
     /// to the real time and execute rounds on the subnets.

rs/pocket_ic_server/tests/test.rs

@@ -1,20 +1,23 @@
-use candid::{Encode, Principal};
+use candid::{Decode, Encode, Principal};
 use ic_agent::agent::CallResponse;
 use ic_cdk::api::management_canister::main::CanisterIdRecord;
 use ic_cdk::api::management_canister::provisional::ProvisionalCreateCanisterWithCyclesArgument;
 use ic_interfaces_registry::{
     RegistryDataProvider, RegistryVersionedRecord, ZERO_REGISTRY_VERSION,
 };
-use ic_management_canister_types::ProvisionalCreateCanisterWithCyclesArgs;
+use ic_management_canister_types::{
+    NodeMetricsHistoryArgs, NodeMetricsHistoryResponse, ProvisionalCreateCanisterWithCyclesArgs,
+};
 use ic_registry_proto_data_provider::ProtoRegistryDataProvider;
 use ic_registry_transport::pb::v1::{
     registry_mutation::Type, RegistryAtomicMutateRequest, RegistryMutation,
 };
 use ic_utils::interfaces::ManagementCanister;
 use nix::sys::signal::Signal;
 use pocket_ic::common::rest::{
-    CreateHttpGatewayResponse, HttpGatewayBackend, HttpGatewayConfig, HttpGatewayDetails,
-    HttpsConfig, InstanceConfig, SubnetConfigSet, SubnetKind, Topology,
+    BlockMakerConfigs, CreateHttpGatewayResponse, HttpGatewayBackend, HttpGatewayConfig,
+    HttpGatewayDetails, HttpsConfig, InstanceConfig, RawSubnetBlockmakerMetrics, SubnetConfigSet,
+    SubnetKind, TickConfigs, Topology,
 };
 use pocket_ic::{update_candid, PocketIc, PocketIcBuilder};
 use rcgen::{CertificateParams, KeyPair};
@@ -1055,6 +1058,65 @@ fn test_query_stats_live() {
     })
 }
 
+#[test]
+fn test_custom_blockmaker() {
+    // Setup PocketIC with initial subnets and records
+    let pocket_ic = PocketIcBuilder::new().with_application_subnet().build();
+    let topology = pocket_ic.topology();
+    let application_subnet = topology.get_app_subnets()[0];
+    let blockmaker_node = topology
+        .subnet_configs
+        .get(&application_subnet)
+        .unwrap()
+        .node_ids[0]
+        .clone();
+
+    let subnets_blockmakers = vec![RawSubnetBlockmakerMetrics {
+        subnet: application_subnet.into(),
+        blockmaker: blockmaker_node,
+        failed_blockmakers: vec![],
+    }];
+
+    let tick_configs = TickConfigs {
+        blockmakers: Some(BlockMakerConfigs {
+            blockmakers_per_subnet: subnets_blockmakers,
+        }),
+    };
+
+    pocket_ic.tick_with_configs(tick_configs.clone());
+    pocket_ic.tick_with_configs(tick_configs.clone());
+    pocket_ic.tick_with_configs(tick_configs.clone());
+    pocket_ic.tick_with_configs(tick_configs.clone());
+    pocket_ic.tick_with_configs(tick_configs.clone());
+
+    pocket_ic
+        //go to next day it should have recorder the metrics
+        .advance_time(std::time::Duration::from_secs(60 * 60 * 25));
+
+    pocket_ic.tick();
+    pocket_ic.tick();
+    pocket_ic.tick();
+    pocket_ic.tick();
+    pocket_ic.tick();
+
+    let response = pocket_ic
+        .query_call(
+            Principal::management_canister(),
+            Principal::anonymous(),
+            "node_metrics_history",
+            Encode!(&NodeMetricsHistoryArgs {
+                subnet_id: application_subnet.into(),
+                start_at_timestamp_nanos: 0,
+            })
+            .unwrap(),
+        )
+        .unwrap();
+
+    let res = Decode!(&response, Vec<NodeMetricsHistoryResponse>).unwrap();
+    println!("{:?}", res);
+    assert!(false)
+}
+
 /// Tests subnet read state requests.
 #[test]
 fn test_subnet_read_state() {

rs/state_machine_tests/tests/multi_subnet.rs

@@ -211,7 +211,7 @@ fn counter_canister_call_test() {
                 .build(),
         )
         .unwrap();
-    env2.execute_round();
+    env2.execute_round(None);
     let wasm_result = env2.await_ingress(msg2_id, MAX_TICKS).unwrap();
     match wasm_result {
         WasmResult::Reply(bytes) => assert_eq!(bytes, vec![123; 2000000]),
@@ -360,7 +360,7 @@ fn counter_canister_call_test() {
     // to induct the ingress message with large payload
     // and all three inter-canister calls with large arguments
     // from the 1st subnet.
-    env2.execute_round();
+    env2.execute_round(None);
     assert!(matches!(
         env2.ingress_status(&msg20_id),
         IngressStatus::Known { .. }
@@ -395,7 +395,7 @@ fn counter_canister_call_test() {
     // This time, we also need to execute one more round on the 2nd subnet
     // to process the response callback of the inter-canister call
     // to the 1st subnet.
-    env2.execute_round();
+    env2.execute_round(None);
 
     let wasm_result = env2.await_ingress(msg20_id, MAX_TICKS).unwrap();
     match wasm_result {