Skip to content

Commit

Permalink
Mock-builder with generic methods support (#1321)
Browse files Browse the repository at this point in the history
* support generic methods for input parameters

* support generic methods for output parameters

* fix nested generic

* minor change

* minor minor change

* fix clippy issues

* apply William suggestions
  • Loading branch information
lemunozm committed Apr 18, 2023
1 parent a564e8f commit 85dd12a
Show file tree
Hide file tree
Showing 5 changed files with 381 additions and 155 deletions.
4 changes: 3 additions & 1 deletion libs/mock-builder/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,11 @@ version = "0.0.1"
[package.metadata.docs.rs]
targets = ["x86_64-unknown-linux-gnu"]

[dependencies]
frame-support = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }

[dev-dependencies]
codec = { package = "parity-scale-codec", version = "3.0.0", features = ["derive"] }
frame-support = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }
frame-system = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }
scale-info = { version = "2.3.0", features = ["derive"] }
sp-core = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }
Expand Down
201 changes: 48 additions & 153 deletions libs/mock-builder/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,177 +39,72 @@
//!
//! Take a look to the [pallet tests](`tests/pallet.rs`) to have a user view of how to use this crate.

/// Provide methods for register/execute calls
pub mod storage {
use std::{any::Any, cell::RefCell, collections::HashMap};
/// Provide functions for register/execute calls
pub mod storage;

/// Identify a call in the call storage
pub type CallId = u64;

trait Callable {
fn as_any(&self) -> &dyn Any;
}

thread_local! {
static CALLS: RefCell<HashMap<CallId, Box<dyn Callable>>>
= RefCell::new(HashMap::default());
}

struct CallWrapper<Input, Output>(Box<dyn Fn(Input) -> Output>);

impl<Input: 'static, Output: 'static> Callable for CallWrapper<Input, Output> {
fn as_any(&self) -> &dyn Any {
self
}
}

/// Register a call into the call storage.
/// The registered call can be uniquely identified by the returned `CallId`.
pub fn register_call<F: Fn(Args) -> R + 'static, Args: 'static, R: 'static>(f: F) -> CallId {
CALLS.with(|state| {
let registry = &mut *state.borrow_mut();
let call_id = registry.len() as u64;
registry.insert(call_id, Box::new(CallWrapper(Box::new(f))));
call_id
})
}

/// Execute a call from the call storage identified by a `call_id`.
pub fn execute_call<Args: 'static, R: 'static>(call_id: CallId, args: Args) -> R {
CALLS.with(|state| {
let registry = &*state.borrow();
let call = registry.get(&call_id).unwrap();
call.as_any()
.downcast_ref::<CallWrapper<Args, R>>()
.expect("Bad mock implementation: expected other function type")
.0(args)
})
}
}
/// Provide functions for handle fuction locations
pub mod location;

use frame_support::{Blake2_128, StorageHasher, StorageMap};
use location::FunctionLocation;
pub use storage::CallId;

/// Prefix that the register functions should have.
pub const MOCK_FN_PREFIX: &str = "mock_";

/// Gives the absolute string identification of a function.
#[macro_export]
macro_rules! function_locator {
() => {{
// Aux function to extract the path
fn f() {}

fn type_name_of<T>(_: T) -> &'static str {
std::any::type_name::<T>()
}
let name = type_name_of(f);
&name[..name.len() - "::f".len()]
}};
/// Register a mock function into the mock function storage.
/// This function should be called with a locator used as a function identification.
pub fn register<Map, L, F, I, O>(locator: L, f: F)
where
Map: StorageMap<<Blake2_128 as StorageHasher>::Output, CallId>,
L: Fn(),
F: Fn(I) -> O + 'static,
I: 'static,
O: 'static,
{
let location = FunctionLocation::from(locator)
.normalize()
.strip_name_prefix(MOCK_FN_PREFIX)
.append_type_signature::<I, O>();

Map::insert(location.hash::<Blake2_128>(), storage::register_call(f));
}

/// Gives the string identification of a function.
/// The identification will be the same no matter if it belongs to a trait or has an `except_`
/// prefix name.
#[macro_export]
macro_rules! call_locator {
() => {{
let path_name = $crate::function_locator!();
let (path, name) = path_name.rsplit_once("::").expect("always ::");

let base_name = name.strip_prefix($crate::MOCK_FN_PREFIX).unwrap_or(name);
let correct_path = path
.strip_prefix("<")
.map(|trait_path| trait_path.split_once(" as").expect("always ' as'").0)
.unwrap_or(path);

format!("{}::{}", correct_path, base_name)
}};
/// Execute a function from the function storage.
/// This function should be called with a locator used as a function identification.
pub fn execute<Map, L, I, O>(locator: L, input: I) -> O
where
Map: StorageMap<<Blake2_128 as StorageHasher>::Output, CallId>,
L: Fn(),
I: 'static,
O: 'static,
{
let location = FunctionLocation::from(locator)
.normalize()
.append_type_signature::<I, O>();

let call_id = Map::try_get(location.hash::<Blake2_128>())
.unwrap_or_else(|_| panic!("Mock was not found. Location: {location:?}"));

storage::execute_call(call_id, input).unwrap_or_else(|err| {
panic!("{err}. Location: {location:?}");
})
}

/// Register a call into the call storage.
/// This macro should be called from the method that wants to register `f`.
/// This macro must be called from a pallet with the `CallIds` storage.
/// Check the main documentation.
/// Register a mock function into the mock function storage.
/// Same as `register()` but with using the locator who calls this macro.
#[macro_export]
macro_rules! register_call {
($f:expr) => {{
use frame_support::StorageHasher;

let call_id = frame_support::Blake2_128::hash($crate::call_locator!().as_bytes());

CallIds::<T>::insert(call_id, $crate::storage::register_call($f));
$crate::register::<CallIds<T>, _, _, _, _>(|| (), $f)
}};
}

/// Execute a call from the call storage.
/// This macro should be called from the method that wants to execute `f`.
/// This macro must be called from a pallet with the `CallIds` storage.
/// Check the main documentation.
/// Execute a function from the function storage.
/// Same as `execute()` but with using the locator who calls this macro.
#[macro_export]
macro_rules! execute_call {
($params:expr) => {{
use frame_support::StorageHasher;

let hash = frame_support::Blake2_128::hash($crate::call_locator!().as_bytes());
let call_id = CallIds::<T>::get(hash).expect(&format!(
"Called to {}, but mock was not found",
$crate::call_locator!()
));

$crate::storage::execute_call(call_id, $params)
($input:expr) => {{
$crate::execute::<CallIds<T>, _, _, _>(|| (), $input)
}};
}

#[cfg(test)]
mod tests {
trait TraitExample {
fn function_locator() -> String;
fn call_locator() -> String;
}

struct Example;

impl Example {
fn mock_function_locator() -> String {
function_locator!().into()
}

fn mock_call_locator() -> String {
call_locator!().into()
}
}

impl TraitExample for Example {
fn function_locator() -> String {
function_locator!().into()
}

fn call_locator() -> String {
call_locator!().into()
}
}

#[test]
fn function_locator() {
assert_eq!(
Example::mock_function_locator(),
"mock_builder::tests::Example::mock_function_locator"
);

assert_eq!(
Example::function_locator(),
"<mock_builder::tests::Example as \
mock_builder::tests::TraitExample>::function_locator"
);
}

#[test]
fn call_locator() {
assert_eq!(
Example::call_locator(),
"mock_builder::tests::Example::call_locator"
);

assert_eq!(Example::call_locator(), Example::mock_call_locator());
}
}
160 changes: 160 additions & 0 deletions libs/mock-builder/src/location.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
use frame_support::StorageHasher;

/// Absolute string identification of function.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct FunctionLocation(String);

impl FunctionLocation {
/// Creates a location for the function which created the given closure used as a locator
pub fn from<F: Fn()>(_: F) -> Self {
let location = std::any::type_name::<F>();
let location = &location[..location.len() - "::{{closure}}".len()];

// Remove generic attributes from signature if it has any
let location = location
.ends_with('>')
.then(|| {
let mut count = 0;
for (i, c) in location.chars().rev().enumerate() {
if c == '>' {
count += 1;
} else if c == '<' {
count -= 1;
if count == 0 {
return location.split_at(location.len() - i - 1).0;
}
}
}
panic!("Expected '<' symbol to close '>'");
})
.unwrap_or(location);

Self(location.into())
}

/// Normalize the location, allowing to identify the function
/// no matter if it belongs to a trait or not.
pub fn normalize(self) -> Self {
let (path, name) = self.0.rsplit_once("::").expect("always ::");
let path = path
.strip_prefix('<')
.map(|trait_path| trait_path.split_once(" as").expect("always ' as'").0)
.unwrap_or(path);

Self(format!("{}::{}", path, name))
}

/// Remove the prefix from the function name.
pub fn strip_name_prefix(self, prefix: &str) -> Self {
let (path, name) = self.0.rsplit_once("::").expect("always ::");
let name = name.strip_prefix(prefix).unwrap_or_else(|| {
panic!(
"Function '{name}' should have a '{prefix}' prefix. Location: {}",
self.0
)
});

Self(format!("{}::{}", path, name))
}

/// Add a representation of the function input and output types
pub fn append_type_signature<I, O>(self) -> Self {
Self(format!(
"{}:{}->{}",
self.0,
std::any::type_name::<I>(),
std::any::type_name::<O>(),
))
}

/// Generate a hash of the location
pub fn hash<Hasher: StorageHasher>(&self) -> Hasher::Output {
Hasher::hash(self.0.as_bytes())
}
}

#[cfg(test)]
mod tests {
use super::*;

const PREFIX: &str = "mock_builder::location::tests";

trait TraitExample {
fn method() -> FunctionLocation;
fn generic_method<A: Into<i32>>(_: impl Into<u32>) -> FunctionLocation;
}

struct Example;

impl Example {
fn mock_method() -> FunctionLocation {
FunctionLocation::from(|| ())
}

fn mock_generic_method<A: Into<i32>>(_: impl Into<u32>) -> FunctionLocation {
FunctionLocation::from(|| ())
}
}

impl TraitExample for Example {
fn method() -> FunctionLocation {
FunctionLocation::from(|| ())
}

fn generic_method<A: Into<i32>>(_: impl Into<u32>) -> FunctionLocation {
FunctionLocation::from(|| ())
}
}

#[test]
fn function_location() {
assert_eq!(
Example::mock_method().0,
format!("{PREFIX}::Example::mock_method")
);

assert_eq!(
Example::mock_generic_method::<i8>(0u8).0,
format!("{PREFIX}::Example::mock_generic_method")
);

assert_eq!(
Example::method().0,
format!("<{PREFIX}::Example as {PREFIX}::TraitExample>::method")
);

assert_eq!(
Example::generic_method::<i8>(0u8).0,
format!("<{PREFIX}::Example as {PREFIX}::TraitExample>::generic_method")
);
}

#[test]
fn normalized_function_location() {
assert_eq!(
Example::mock_method().normalize().0,
format!("{PREFIX}::Example::mock_method")
);

assert_eq!(
Example::method().normalize().0,
format!("{PREFIX}::Example::method")
);
}

#[test]
fn striped_function_location() {
assert_eq!(
Example::mock_method().strip_name_prefix("mock_").0,
format!("{PREFIX}::Example::method")
);
}

#[test]
fn appended_type_signature() {
assert_eq!(
Example::mock_method().append_type_signature::<i8, u8>().0,
format!("{PREFIX}::Example::mock_method:i8->u8")
);
}
}
Loading

0 comments on commit 85dd12a

Please sign in to comment.