chore(boundary): Add pocket-ic integration tests for rate-limit canis…

…ter (#2360)

To run this test manually, use:

```
bazel test --config=lint //rs/boundary_node/rate_limits/integration_tests/...
```

---------

Co-authored-by: IDX GitLab Automation <[email protected]>
Co-authored-by: Nikolay Komarevskiy <[email protected]>

rs/boundary_node/rate_limits/Cargo.toml

@@ -27,6 +27,9 @@ sha2 = { workspace = true }
 thiserror = { workspace = true }
 uuid = { workspace = true, features = ['serde', 'v4'] }
 
+[dev-dependencies]
+rate-limit-canister-integration-tests = { path = "./integration_tests" }
+
 [lib]
 crate-type = ["cdylib"]
 path = "canister/lib.rs"

rs/boundary_node/rate_limits/api/src/lib.rs

@@ -68,7 +68,7 @@ pub struct OutputRuleMetadata {
     pub removed_in_version: Option<Version>,
 }
 
-#[derive(CandidType, Deserialize, Debug)]
+#[derive(CandidType, Deserialize, Debug, Clone)]
 pub struct InitArg {
     pub registry_polling_period_secs: u64,
     pub authorized_principal: Option<Principal>,

rs/boundary_node/rate_limits/integration_tests/BUILD.bazel

@@ -0,0 +1,74 @@
+load("@rules_rust//rust:defs.bzl", "rust_library")
+load("//bazel:defs.bzl", "rust_test_suite_with_extra_srcs")
+
+package(default_visibility = ["//visibility:public"])
+
+DEPENDENCIES = [
+    # Keep sorted.
+    "//rs/boundary_node/rate_limits/api:rate_limits_api",
+    "//rs/nervous_system/integration_tests:nervous_system_integration_tests",
+    "//rs/types/base_types",
+    "@crate_index//:assert_matches",
+    "@crate_index//:candid",
+    "@crate_index//:serde",
+    "@crate_index//:tokio",
+] + select({
+    "@rules_rust//rust/platform:wasm32-unknown-unknown": [],
+    "//conditions:default": [
+        "//packages/pocket-ic",
+        "//rs/crypto/sha2",
+        "//rs/nns/constants",
+        "//rs/registry/keys",
+        "//rs/registry/transport",
+        "//rs/rust_canisters/canister_test",
+        "//rs/nns/test_utils:test_utils",
+    ],
+})
+
+MACRO_DEPENDENCIES = []
+
+DEV_DEPENDENCIES = []
+
+MACRO_DEV_DEPENDENCIES = []
+
+ALIASES = {}
+
+DEV_DATA = [
+    "@mainnet_nns_registry_canister//file",
+    "//rs/registry/canister:registry-canister",
+    "//rs/pocket_ic_server:pocket-ic-server",
+    "//rs/boundary_node/rate_limits:rate_limit_canister",
+]
+
+DEV_ENV = {
+    "CARGO_MANIFEST_DIR": "rs/nns/integration_tests",
+    "REGISTRY_CANISTER_WASM_PATH": "$(rootpath //rs/registry/canister:registry-canister)",
+    "MAINNET_REGISTRY_CANISTER_WASM_PATH": "$(rootpath @mainnet_nns_registry_canister//file)",
+    "POCKET_IC_BIN": "$(rootpath //rs/pocket_ic_server:pocket-ic-server)",
+    "RATE_LIMITS_CANISTER_WASM_PATH": "$(rootpath //rs/boundary_node/rate_limits:rate_limit_canister)",
+}
+
+rust_library(
+    name = "rate_limit_canister_integration_tests",
+    testonly = True,
+    srcs = glob(["src/**/*.rs"]),
+    aliases = ALIASES,
+    crate_name = "rate_limit_canister_integration_tests",
+    proc_macro_deps = MACRO_DEPENDENCIES,
+    deps = DEPENDENCIES,
+)
+
+rust_test_suite_with_extra_srcs(
+    name = "integration_tests_test",
+    srcs = glob(
+        ["tests/**/*.rs"],
+    ),
+    aliases = ALIASES,
+    data = DEV_DATA,
+    env = DEV_ENV,
+    extra_srcs = [],
+    flaky = False,
+    proc_macro_deps = MACRO_DEPENDENCIES + MACRO_DEV_DEPENDENCIES,
+    tags = [],
+    deps = [":rate_limit_canister_integration_tests"] + DEPENDENCIES + DEV_DEPENDENCIES,
+)

rs/boundary_node/rate_limits/integration_tests/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "rate-limit-canister-integration-tests"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+description.workspace = true
+documentation.workspace = true
+
+[dependencies]
+rate-limits-api = { path = "../api" }
+ic-nervous-system-integration-tests = { path = "../../../nervous_system/integration_tests" }
+ic-base-types = { path = "../../../types/base_types" }
+assert_matches = { workspace = true }
+candid = { workspace = true }
+serde = { workspace = true }
+tokio = { workspace = true }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+pocket-ic = { path = "../../../../packages/pocket-ic" }
+ic-crypto-sha2 = { path = "../../../crypto/sha2" }
+ic-nns-constants = { path = "../../../nns/constants" }
+ic-registry-keys = { path = "../../../registry/keys" }
+ic-registry-transport = { path = "../../../registry/transport" }
+canister-test = { path = "../../../rust_canisters/canister_test" }
+ic-nns-test-utils = { path = "../../../nns/test_utils" }
+

rs/boundary_node/rate_limits/integration_tests/src/lib.rs

@@ -0,0 +1,14 @@
+//! Integration tests for the rate-limit canister.
+//!
+//! Each test creates a PocketIc instance, installs the registry and rate-limit canisters, and then
+//! proceeds to perform operations and verify they completed successfully, and
+//! that the state is the expected one. State inspection is done via the public
+//! canister API.
+//!
+//! This is not a library at all. However, if this was under `tests/`, then each
+//! file would become its own crate, and the tests would run sequentially. By
+//! pretending it's a library with several modules inside, `cargo test` is
+//! supposed to run all tests in parallel, because they are all in the same
+//! crate.
+
+pub mod pocket_ic_helpers;

rs/boundary_node/rate_limits/integration_tests/src/pocket_ic_helpers.rs

@@ -0,0 +1,159 @@
+use candid::{CandidType, Decode, Encode, Principal};
+use canister_test::Project;
+use canister_test::Wasm;
+use ic_crypto_sha2::Sha256;
+use ic_nns_constants::{REGISTRY_CANISTER_ID, ROOT_CANISTER_ID};
+use ic_nns_test_utils::common::{
+    build_mainnet_registry_wasm, build_registry_wasm, NnsInitPayloadsBuilder,
+};
+use ic_registry_transport::pb::v1::RegistryAtomicMutateRequest;
+use pocket_ic::{nonblocking::PocketIc, CanisterSettings, PocketIcBuilder, WasmResult};
+use rate_limits_api::InitArg;
+use serde::de::DeserializeOwned;
+
+/// Builds the WASM for the rate-limit canister.
+pub fn build_rate_limits_wasm() -> Wasm {
+    Project::cargo_bin_maybe_from_env("rate-limits-canister", &[])
+}
+
+pub async fn install_registry_canister(
+    pocket_ic: &PocketIc,
+    with_mainnet_nns_canister_versions: bool,
+    custom_initial_registry_mutations: Option<Vec<RegistryAtomicMutateRequest>>,
+) {
+    let mut nns_init_payload_builder = NnsInitPayloadsBuilder::new();
+
+    if let Some(custom_initial_registry_mutations) = custom_initial_registry_mutations {
+        nns_init_payload_builder.with_initial_mutations(custom_initial_registry_mutations);
+    } else {
+        nns_init_payload_builder.with_initial_invariant_compliant_mutations();
+    }
+
+    let nns_init_payload = nns_init_payload_builder.build();
+
+    let registry_wasm = if with_mainnet_nns_canister_versions {
+        build_mainnet_registry_wasm()
+    } else {
+        build_registry_wasm()
+    };
+
+    ic_nervous_system_integration_tests::pocket_ic_helpers::install_canister(
+        pocket_ic,
+        "registry canister",
+        REGISTRY_CANISTER_ID,
+        Encode!(&nns_init_payload.registry).unwrap(),
+        registry_wasm,
+        Some(ROOT_CANISTER_ID.get()),
+    )
+    .await;
+}
+
+pub async fn install_canister(
+    pocket_ic: &PocketIc,
+    canister_name: &str,
+    subnet_id: Principal,
+    arg: Vec<u8>,
+    wasm: Wasm,
+) -> Principal {
+    let memory_allocation = None;
+    let controllers = None;
+    let sender = None;
+
+    let settings = Some(CanisterSettings {
+        memory_allocation,
+        controllers,
+        ..Default::default()
+    });
+
+    let canister_id = pocket_ic
+        .create_canister_on_subnet(sender, settings, subnet_id)
+        .await;
+
+    pocket_ic
+        .install_canister(canister_id, wasm.bytes(), arg, sender)
+        .await;
+
+    println!(
+        "Installed {canister_name} with canister_id = {canister_id} on subnet_id = {subnet_id}",
+    );
+
+    canister_id
+}
+
+pub async fn get_installed_wasm_hash(pocket_ic: &PocketIc, canister_id: Principal) -> [u8; 32] {
+    let module_hash = pocket_ic
+        .canister_status(canister_id, None)
+        .await
+        .unwrap()
+        .module_hash
+        .unwrap();
+
+    module_hash.try_into().unwrap_or_else(|v: Vec<_>| {
+        panic!("Expected a Vec of length 32 but it has {} bytes.", v.len())
+    })
+}
+
+pub async fn canister_call<R: DeserializeOwned + CandidType>(
+    pocket_ic: &PocketIc,
+    method: &str,
+    canister_id: Principal,
+    sender: Principal,
+    payload: Vec<u8>,
+) -> Result<R, String> {
+    let result = pocket_ic
+        .update_call(canister_id, sender, method, payload)
+        .await
+        .map_err(|err| err.to_string())?;
+
+    let result = match result {
+        WasmResult::Reply(result) => result,
+        WasmResult::Reject(s) => panic!("Call to add_config failed: {:#?}", s),
+    };
+
+    let decoded: R = Decode!(&result, R).unwrap();
+
+    Ok(decoded)
+}
+
+pub async fn setup_subnets_and_registry_canister() -> PocketIc {
+    let pocket_ic = PocketIcBuilder::new()
+        .with_nns_subnet()
+        .with_ii_subnet()
+        .build_async()
+        .await;
+
+    // Install registry canister. It is the only canister that rate-limit canister interacts with.
+    let with_mainnet_nns_canister_versions = false;
+    install_registry_canister(&pocket_ic, with_mainnet_nns_canister_versions, None).await;
+
+    pocket_ic
+}
+
+pub async fn install_rate_limit_canister_on_ii_subnet(
+    pocket_ic: &PocketIc,
+    init_arg: InitArg,
+) -> (Principal, Wasm) {
+    let wasm = build_rate_limits_wasm();
+    let wasm_hash = Sha256::hash(&wasm.clone().bytes());
+
+    let ii_subnet_id = {
+        let topology = pocket_ic.topology().await;
+        topology.get_ii().unwrap()
+    };
+
+    let canister_id = install_canister(
+        pocket_ic,
+        "rate-limit canister",
+        ii_subnet_id,
+        Encode!(&init_arg).unwrap(),
+        wasm.clone(),
+    )
+    .await;
+
+    assert_eq!(
+        get_installed_wasm_hash(pocket_ic, canister_id).await,
+        wasm_hash,
+    );
+
+    (canister_id, wasm)
+}