[WIP] reference impl of mocked corss-contract calls in ink unit tests

Cargo.toml

@@ -72,6 +72,7 @@ pausable = ["openbrush_contracts/pausable"]
 timelock_controller = ["openbrush_contracts/timelock_controller"]
 proxy = ["openbrush_contracts/proxy"]
 diamond = ["openbrush_contracts/diamond"]
+mockable = ["openbrush_lang/mockable"]
 
 [profile.release]
 panic = "abort"

lang/Cargo.toml

@@ -43,4 +43,5 @@ std = [
     "ink_metadata/std",
     "scale/std",
     "scale-info/std",
-]
+]
+mockable = []

lang/codegen/src/trait_definition.rs

@@ -48,6 +48,7 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
         return (quote! {}).into()
     }
     let attrs: proc_macro2::TokenStream = _attrs.into();
+    let (mock_type, attrs) = extract_mock_config(attrs);
     let mut trait_item: ItemTrait = parse2(_input).unwrap();
     let trait_without_ink_attrs;
     let ink_code;
@@ -60,6 +61,7 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
         }
     });
 
+    let mut maybe_use_mock_env = quote! {};
     if contains_ink.is_some() {
         add_selectors_attribute(&mut trait_item);
         // Brackets to force the unlock of the file after the update of the trait definition
@@ -97,7 +99,7 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
             }
         });
 
-        let wrapper_trait = generate_wrapper(ink_trait.clone());
+        let wrapper_trait = generate_wrapper(ink_trait.clone(), mock_type.clone());
 
         ink_code = quote! {
             #[allow(non_camel_case_types)]
@@ -113,6 +115,14 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
                 #ink_trait
             }
         };
+
+        let pub_mock_env_ident = format_ident!("mock_{}", trait_item.ident.to_string().to_lowercase());
+        maybe_use_mock_env = quote! {
+            #[cfg(any(test, feature = "mockable"))]
+            pub mod #pub_mock_env_ident {
+                pub use super :: #namespace_ident :: { mock_env as env , using , deploy };
+            }
+        };
     } else {
         trait_without_ink_attrs = trait_item;
         ink_code = quote! {};
@@ -124,6 +134,8 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
         #trait_without_ink_attrs
 
         #ink_code
+
+        #maybe_use_mock_env
     };
     code.into()
 }
@@ -167,11 +179,12 @@ fn transform_to_ink_trait(mut trait_item: ItemTrait) -> ItemTrait {
     trait_item
 }
 
-fn generate_wrapper(ink_trait: ItemTrait) -> proc_macro2::TokenStream {
+fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> proc_macro2::TokenStream {
     let trait_ident = ink_trait.ident.clone();
     let trait_wrapper_ident = format_ident!("{}Wrapper", ink_trait.ident);
     let mut def_messages = vec![];
     let mut impl_messages = vec![];
+    let mock_address_pattern = name_to_raw_account(&format!("Mock{}", ink_trait.ident));
     ink_trait
         .items
         .clone()
@@ -252,15 +265,37 @@ fn generate_wrapper(ink_trait: ItemTrait) -> proc_macro2::TokenStream {
                 >;
             });
 
+            let message_test_impl = match &mock_type {
+                Some(_mock_ty) => quote! {
+                    mock_env :: with(|ctx| {
+                        let mut mock_ref = ctx.register.get_mut(self).expect("not an address of mocked contract");
+                        ctx.stack.push(&self);
+                        let result = mock_ref.borrow_mut(). #message_ident (
+                            #( #input_bindings , )*
+                        );
+                        ctx.stack.pop();
+                        result
+                    }).expect("mock object not set")
+                },
+                None => quote! { ::core::panic!("cross-contract call is not supported in ink tests; try to set a mock object?") }
+            };
+
             impl_messages.push(quote! {
                 #[inline]
                 fn #message_ident(
                     & self
                     #( , #input_bindings : #input_types )*
                 ) -> #output_ty {
-                    Self::#message_builder_ident(self #( , #input_bindings)*)
-                        .fire()
-                        .unwrap_or_else(|err| ::core::panic!("{}: {:?}", #panic_str, err))
+                    #[cfg(not(any(test, feature = "mockable")))]
+                    {
+                        Self::#message_builder_ident(self #( , #input_bindings)*)
+                            .fire()
+                            .unwrap_or_else(|err| ::core::panic!("{}: {:?}", #panic_str, err))
+                    }
+                    #[cfg(any(test, feature = "mockable"))]
+                    {
+                        #message_test_impl
+                    }
                 }
 
                 #[inline]
@@ -292,6 +327,62 @@ fn generate_wrapper(ink_trait: ItemTrait) -> proc_macro2::TokenStream {
     let impl_messages = impl_messages.iter();
     let def_messages = def_messages.iter();
 
+    let maybe_mock_environmental = match mock_type {
+        Some(ty) => {
+            quote! {
+                #[cfg(any(test, feature = "mockable"))]
+                pub struct Context {
+                    pub stack: ::openbrush::traits::mock::SharedCallStack,
+                    pub register: std::collections::BTreeMap<
+                        ::openbrush::traits::AccountId,
+                        std::rc::Rc<std::cell::RefCell< #ty >>
+                    >
+                }
+
+                #[cfg(any(test, feature = "mockable"))]
+                ::environmental::environmental!(
+                    pub mock_env : Context
+                );
+
+                #[cfg(any(test, feature = "mockable"))]
+                pub fn using<F: FnOnce()>(
+                    stack: ::openbrush::traits::mock::SharedCallStack,
+                    f: F
+                ) {
+                    let mut env = Context {
+                        stack,
+                        register: Default::default()
+                    };
+                    mock_env::using(&mut env, f);
+                }
+
+                #[cfg(any(test, feature = "mockable"))]
+                pub fn deploy(inner_contract : #ty) -> (::openbrush::traits::mock::Addressable< #ty >) {
+                    let contract: std::rc::Rc<std::cell::RefCell< #ty >> = std::rc::Rc::new(
+                        std::cell::RefCell::< #ty >::new(inner_contract)
+                    );
+                    let (account_id, contract, stack) = mock_env::with(|ctx| {
+                        let n: u8 = ctx.register.len().try_into()
+                            .expect("too many contracts to fit into u8");
+                        let mut pat = [ #( #mock_address_pattern,  )* ];
+                        pat[31] = n;
+                        let account_id: ::openbrush::traits::AccountId = pat.into();
+
+                        ctx.register.insert(account_id.clone(), contract.clone());
+                        (account_id, contract, ctx.stack.clone())
+                    }).expect("must call within `using()`");
+
+                    ::openbrush::traits::mock::Addressable::new(
+                        account_id,
+                        contract,
+                        stack,
+                    )
+                }
+            }
+        }
+        None => quote! {},
+    };
+
     quote! {
         pub trait #trait_wrapper_ident {
             #( #def_messages )*
@@ -300,6 +391,8 @@ fn generate_wrapper(ink_trait: ItemTrait) -> proc_macro2::TokenStream {
         impl #trait_wrapper_ident for ::openbrush::traits::AccountId {
             #( #impl_messages )*
         }
+
+        #maybe_mock_environmental
     }
 }
 
@@ -334,3 +427,59 @@ fn remove_ink_attrs(mut trait_item: ItemTrait) -> ItemTrait {
     });
     trait_item
 }
+
+/// Extracts the mocking related macro args out from the input
+///
+/// Return a tuple of an optional mock target and the args without the mock target
+fn extract_mock_config(attr: TokenStream) -> (Option<TokenStream>, TokenStream) {
+    let attr_args = syn::parse2::<attr_args::AttributeArgs>(attr).expect("unable to parse trait_definition attribute");
+
+    let (mock_args, ink_args): (Vec<_>, Vec<_>) = attr_args.into_iter().partition(|arg| arg.name.is_ident("mock"));
+
+    let mock_type = mock_args.first().map(|mock_attr| {
+        let ty = &mock_attr.value;
+        quote! { #ty }
+    });
+    let ink_attrs = quote! {
+        #( #ink_args , ) *
+    };
+    (mock_type, ink_attrs)
+}
+
+/// Returns a `[u8; 32]` filled with the give str with zero padding.
+fn name_to_raw_account(name: &str) -> [u8; 32] {
+    let mut v = name.as_bytes().to_vec();
+    v.resize(32, 0);
+    v.try_into().expect("length is 32; qed.")
+}
+
+mod attr_args;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn macro_works() {
+        let r = generate(
+            quote! {
+                mock = MyMockType,
+                namespace = ::name::space
+            },
+            quote! {
+                pub trait SubmittableOracle {
+                    #[ink(message)]
+                    fn admin(&self) -> AccountId;
+
+                    #[ink(message)]
+                    fn verifier(&self) -> Verifier;
+
+                    #[ink(message)]
+                    fn attest(&self, arg: String) -> Result<Attestation, ()>;
+                }
+            },
+        );
+
+        println!("OUTPUT:\n\n{:}", r);
+    }
+}