Skip to content

Commit

Permalink
Merge branch 'main' into revert-202-EPROD-653-define-how-to-handle-pa…
Browse files Browse the repository at this point in the history
…nics-in-the-task-scheduler
  • Loading branch information
ufoscout committed Mar 12, 2024
2 parents cd1b20d + c229d21 commit 6883e7d
Show file tree
Hide file tree
Showing 10 changed files with 294 additions and 16 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -80,5 +80,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"
1 change: 1 addition & 0 deletions ic-canister-client/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
11 changes: 11 additions & 0 deletions ic-canister-client/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions ic-canister-client/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
157 changes: 157 additions & 0 deletions ic-canister-client/src/state_machine_tests.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
2 changes: 2 additions & 0 deletions ic-exports/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand All @@ -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 }
Expand Down
102 changes: 102 additions & 0 deletions ic-exports/src/ic_test_state_machine.rs
Original file line number Diff line number Diff line change
@@ -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())
}
3 changes: 3 additions & 0 deletions ic-exports/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
27 changes: 12 additions & 15 deletions ic-exports/src/pocket_ic/nio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use std::sync::Arc;
use std::time::{Duration, SystemTime};

use candid::Principal;
use ic_cdk::api::management_canister::main::CanisterSettings;
use ic_cdk::api::management_canister::provisional::CanisterId;
use pocket_ic::common::rest::{BlobCompression, BlobId};
use pocket_ic::{CallError, PocketIc, UserError, WasmResult};
Expand Down Expand Up @@ -188,21 +189,17 @@ impl PocketIcAsync {
.unwrap()
}

// todo: implement after pocket-ic will update `ic-cdk` dependency version.
// because currently they use `0.10` and canister-sdk uses `0.11`.
// Issue: https://infinityswap.atlassian.net/browse/EPROD-601
//
/// Create a canister with custom settings.
// pub async fn create_canister_with_settings(
// &self,
// settings: Option<CanisterSettings>,
// sender: Option<Principal>,
// ) -> CanisterId {
// let client = self.0.clone();
// tokio::task::spawn_blocking(move || client.create_canister_with_settings(settings, sender))
// .await
// .unwrap()
// }
// Create a canister with custom settings.
pub async fn create_canister_with_settings(
&self,
settings: Option<CanisterSettings>,
sender: Option<Principal>,
) -> CanisterId {
let client = self.0.clone();
tokio::task::spawn_blocking(move || client.create_canister_with_settings(sender, settings))
.await
.unwrap()
}

/// Install a WASM module on an existing canister.
pub async fn install_canister(
Expand Down
1 change: 0 additions & 1 deletion ic-helpers/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ pub mod utils;
pub use utils::*;

pub mod principal;
pub use principal::*;

pub mod types;
pub use types::*;
Expand Down

0 comments on commit 6883e7d

Please sign in to comment.