Restore ic-state-machine-test client

Cargo.toml

@@ -81,5 +81,6 @@ ic-cdk = "0.12"
 ic-cdk-macros = "0.8"
 ic-cdk-timers = "0.6"
 ic-ledger-types = "0.9"
+ic-test-state-machine-client = "3"
 icrc-ledger-types = "0.1.0"
 pocket-ic = "2.2"

ic-canister-client/Cargo.toml

@@ -10,6 +10,7 @@ description = "Client for interacting with an IC Canister"
 default = []
 ic-agent-client = ["dep:ic-agent"]
 pocket-ic-client = ["dep:tokio", "ic-exports/pocket-ic-tests-async"]
+state-machine-tests-client = ["dep:tokio", "ic-exports/ic-test-state-machine"]
 
 [dependencies]
 async-trait = { workspace = true }

ic-canister-client/src/error.rs

@@ -13,6 +13,10 @@ pub enum CanisterClientError {
     #[error("ic agent error: {0}")]
     IcAgentError(#[from] ic_agent::agent::AgentError),
 
+    #[cfg(feature = "state-machine-tests-client")]
+    #[error("state machine test error: {0}")]
+    StateMachineTestError(ic_exports::ic_test_state_machine::UserError),
+
     #[cfg(feature = "pocket-ic-client")]
     #[error("pocket-ic test error: {0:?}")]
     PocketIcTestError(ic_exports::pocket_ic::CallError),
@@ -32,6 +36,13 @@ impl From<ic_exports::pocket_ic::UserError> for CanisterClientError {
     }
 }
 
+#[cfg(feature = "state-machine-tests-client")]
+impl From<ic_exports::ic_test_state_machine::UserError> for CanisterClientError {
+    fn from(error: ic_exports::ic_test_state_machine::UserError) -> Self {
+        CanisterClientError::StateMachineTestError(error)
+    }
+}
+
 pub type CanisterClientResult<T> = Result<T, CanisterClientError>;
 
 /// This tuple is returned incase of IC errors such as Network, canister error.

ic-canister-client/src/lib.rs

@@ -5,6 +5,9 @@ pub mod client;
 pub mod error;
 pub mod ic_client;
 
+#[cfg(feature = "state-machine-tests-client")]
+pub mod state_machine_tests;
+
 #[cfg(feature = "pocket-ic-client")]
 pub mod pocket_ic;
 
@@ -17,3 +20,5 @@ pub use ic_agent;
 pub use ic_client::IcCanisterClient;
 #[cfg(feature = "pocket-ic-client")]
 pub use pocket_ic::PocketIcClient;
+#[cfg(feature = "state-machine-tests-client")]
+pub use state_machine_tests::StateMachineCanisterClient;

ic-canister-client/src/state_machine_tests.rs

@@ -0,0 +1,157 @@
+use std::sync::Arc;
+
+use candid::utils::ArgumentEncoder;
+use candid::{CandidType, Decode, Principal};
+use ic_exports::ic_kit::RejectionCode;
+use ic_exports::ic_test_state_machine::{StateMachine, WasmResult};
+use serde::Deserialize;
+use tokio::sync::Mutex;
+
+use crate::{CanisterClient, CanisterClientError, CanisterClientResult};
+
+/// A client for interacting with a canister inside dfinity's
+/// state machine tests framework.
+#[derive(Clone)]
+pub struct StateMachineCanisterClient {
+    state_machine: Arc<Mutex<StateMachine>>,
+    canister: Principal,
+    caller: Principal,
+}
+
+impl StateMachineCanisterClient {
+    /// Creates a new instance of a StateMachineCanisterClient.
+    pub fn new(
+        state_machine: Arc<Mutex<StateMachine>>,
+        canister: Principal,
+        caller: Principal,
+    ) -> Self {
+        Self {
+            state_machine,
+            canister,
+            caller,
+        }
+    }
+
+    /// Returns the caller of the canister.
+    pub fn caller(&self) -> Principal {
+        self.caller
+    }
+
+    /// Replace the caller.
+    pub fn set_caller(&mut self, caller: Principal) {
+        self.caller = caller;
+    }
+
+    /// Returns the canister of the canister.
+    pub fn canister(&self) -> Principal {
+        self.canister
+    }
+
+    /// Replace the canister to call.
+    pub fn set_canister(&mut self, canister: Principal) {
+        self.canister = canister;
+    }
+
+    /// Returns the state machine of the canister.
+    pub fn state_machine(&self) -> &Mutex<StateMachine> {
+        self.state_machine.as_ref()
+    }
+
+    /// Performs a blocking action with state machine and awaits the result.
+    ///
+    /// Arguments of the closure `f`:
+    /// 1) `env` - The state machine environment.
+    /// 2) `canister` - The canister principal.
+    /// 3) `caller` - The caller principal.
+    pub async fn with_state_machine<F, R>(&self, f: F) -> R
+    where
+        F: Send + FnOnce(&StateMachine, Principal, Principal) -> R + 'static,
+        R: Send + 'static,
+    {
+        let client = self.state_machine.clone();
+        let cansiter = self.canister;
+        let caller = self.caller;
+
+        tokio::task::spawn_blocking(move || {
+            let locked_client = client.blocking_lock();
+            f(&locked_client, cansiter, caller)
+        })
+        .await
+        .unwrap()
+    }
+
+    pub async fn update<T, R>(&self, method: &str, args: T) -> CanisterClientResult<R>
+    where
+        T: ArgumentEncoder + Send + Sync,
+        R: for<'de> Deserialize<'de> + CandidType,
+    {
+        let args = candid::encode_args(args)?;
+        let method = String::from(method);
+
+        let call_result = self
+            .with_state_machine(move |env, canister, caller| {
+                env.update_call(canister, caller, &method, args)
+            })
+            .await?;
+
+        let reply = match call_result {
+            WasmResult::Reply(reply) => reply,
+            WasmResult::Reject(e) => {
+                return Err(CanisterClientError::CanisterError((
+                    RejectionCode::CanisterError,
+                    e,
+                )));
+            }
+        };
+
+        let decoded = Decode!(&reply, R)?;
+        Ok(decoded)
+    }
+
+    pub async fn query<T, R>(&self, method: &str, args: T) -> CanisterClientResult<R>
+    where
+        T: ArgumentEncoder + Send + Sync,
+        R: for<'de> Deserialize<'de> + CandidType,
+    {
+        let args = candid::encode_args(args)?;
+        let method = String::from(method);
+
+        let call_result = self
+            .with_state_machine(move |env, canister, caller| {
+                env.query_call(canister, caller, &method, args)
+            })
+            .await?;
+
+        let reply = match call_result {
+            WasmResult::Reply(reply) => reply,
+            WasmResult::Reject(e) => {
+                return Err(CanisterClientError::CanisterError((
+                    RejectionCode::CanisterError,
+                    e,
+                )));
+            }
+        };
+
+        let decoded = Decode!(&reply, R)?;
+        Ok(decoded)
+    }
+}
+
+#[async_trait::async_trait]
+impl CanisterClient for StateMachineCanisterClient {
+    async fn update<T, R>(&self, method: &str, args: T) -> CanisterClientResult<R>
+    where
+        T: ArgumentEncoder + Send + Sync,
+        R: for<'de> Deserialize<'de> + CandidType,
+    {
+        StateMachineCanisterClient::update(self, method, args).await
+    }
+
+    async fn query<T, R>(&self, method: &str, args: T) -> CanisterClientResult<R>
+    where
+        T: ArgumentEncoder + Send + Sync,
+        R: for<'de> Deserialize<'de> + CandidType,
+    {
+        StateMachineCanisterClient::query(self, method, args).await
+    }
+}

ic-exports/Cargo.toml

@@ -10,6 +10,7 @@ edition.workspace = true
 default = []
 ledger = ["ic-ledger-types"]
 icrc = ["icrc-ledger-types"]
+ic-test-state-machine = ["flate2", "ic-test-state-machine-client", "log", "once_cell", "reqwest"]
 pocket-ic-tests = ["flate2", "pocket-ic", "log", "once_cell", "reqwest"]
 pocket-ic-tests-async = ["pocket-ic-tests", "tokio"]
 
@@ -22,6 +23,7 @@ ic-cdk-timers = { workspace = true }
 ic-crypto-getrandom-for-wasm = { path = "../ic-crypto-getrandom-for-wasm" }
 ic-kit = { path = "../ic-kit" }
 ic-ledger-types = { workspace = true, optional = true }
+ic-test-state-machine-client = { workspace = true, optional = true }
 icrc-ledger-types = { workspace = true, optional = true }
 pocket-ic = { workspace = true, optional = true }
 serde = { workspace = true }

ic-exports/src/ic_test_state_machine.rs

@@ -0,0 +1,102 @@
+use std::fs::{create_dir_all, File};
+use std::io::*;
+use std::path::Path;
+use std::time::Duration;
+
+use flate2::read::GzDecoder;
+pub use ic_test_state_machine_client::*;
+use log::*;
+use once_cell::sync::OnceCell;
+
+pub const IC_STATE_MACHINE_BINARY_HASH: &str = "48da85ee6c03e8c15f3e90b21bf9ccae7b753ee6";
+
+/// Returns the path to the ic-test-state-machine binary.
+/// If the binary is not present, it downloads it.
+/// See: https://github.com/dfinity/test-state-machine-client
+///
+/// It supports only linux and macos
+///
+/// The search_path variable is the folder where to search for the binary
+/// or to download it if not present
+pub fn get_ic_test_state_machine_client_path(search_path: &str) -> String {
+    static FILES: OnceCell<String> = OnceCell::new();
+    FILES.get_or_init(|| download_binary(search_path)).clone()
+}
+
+fn download_binary(base_path: &str) -> String {
+    let platform = match std::env::consts::OS {
+        "linux" => "linux",
+        "macos" => "darwin",
+        _ => panic!("ic_test_state_machine_client requires linux or macos"),
+    };
+
+    let output_file_name = "ic-test-state-machine";
+    let gz_file_name = format!("{output_file_name}.gz");
+    let download_url = format!("https://download.dfinity.systems/ic/{IC_STATE_MACHINE_BINARY_HASH}/binaries/x86_64-{platform}/{gz_file_name}");
+
+    let dest_path_name = format!("{}/{}", base_path, "ic_test_state_machine");
+    let dest_dir_path = Path::new(&dest_path_name);
+    let gz_dest_file_path = format!("{}/{}", dest_path_name, gz_file_name);
+    let output_dest_file_path = format!("{}/{}", dest_path_name, output_file_name);
+
+    if !Path::new(&output_dest_file_path).exists() {
+        // Download file
+        {
+            info!(
+                "ic-test-state-machine binarey not found, downloading binary from: {download_url}"
+            );
+
+            let response = reqwest::blocking::Client::builder()
+                .timeout(Duration::from_secs(120))
+                .build()
+                .unwrap()
+                .get(download_url)
+                .send()
+                .unwrap();
+
+            create_dir_all(dest_dir_path).unwrap();
+
+            let mut file = match File::create(&gz_dest_file_path) {
+                Err(why) => panic!("couldn't create {}", why),
+                Ok(file) => file,
+            };
+            let content = response.bytes().unwrap();
+            info!("ic-test-state-machine.gz file length: {}", content.len());
+            file.write_all(&content).unwrap();
+            file.flush().unwrap();
+        }
+
+        // unzip file
+        {
+            info!(
+                "unzip ic-test-state-machine to [{}]",
+                dest_dir_path.to_str().unwrap()
+            );
+            let tar_gz = File::open(gz_dest_file_path).unwrap();
+            let mut tar = GzDecoder::new(tar_gz);
+            let mut temp = vec![];
+            tar.read_to_end(&mut temp).unwrap();
+
+            let mut output = File::create(&output_dest_file_path).unwrap();
+            output.write_all(&temp).unwrap();
+            output.flush().unwrap();
+
+            #[cfg(target_family = "unix")]
+            {
+                use std::os::unix::prelude::PermissionsExt;
+                let mut perms = std::fs::metadata(&output_dest_file_path)
+                    .unwrap()
+                    .permissions();
+                perms.set_mode(0o770);
+                std::fs::set_permissions(&output_dest_file_path, perms).unwrap();
+            }
+        }
+    }
+    output_dest_file_path
+}
+
+#[test]
+fn should_get_ic_test_state_machine_client_path() {
+    let path = get_ic_test_state_machine_client_path("../target");
+    assert!(Path::new(&path).exists())
+}

ic-exports/src/lib.rs

@@ -21,3 +21,6 @@ pub mod icrc_types {
 
 #[cfg(feature = "pocket-ic-tests")]
 pub mod pocket_ic;
+
+#[cfg(feature = "ic-test-state-machine")]
+pub mod ic_test_state_machine;

ic-helpers/src/lib.rs

@@ -4,7 +4,6 @@ pub mod utils;
 pub use utils::*;
 
 pub mod principal;
-pub use principal::*;
 
 pub mod types;
 pub use types::*;