From efe7345ce20264d871c22274ca102ba5d03bc9e2 Mon Sep 17 00:00:00 2001
From: h4x3rotab <h4x3rotab@gmail.com>
Date: Wed, 15 Jun 2022 12:11:50 +0000
Subject: [PATCH 1/5] wip: mvp

---
 lang/codegen/src/trait_definition.rs          | 101 +++++-
 .../codegen/src/trait_definition/attr_args.rs | 307 ++++++++++++++++++
 2 files changed, 403 insertions(+), 5 deletions(-)
 create mode 100644 lang/codegen/src/trait_definition/attr_args.rs

diff --git a/lang/codegen/src/trait_definition.rs b/lang/codegen/src/trait_definition.rs
index 807372dd8..2f39f2a7c 100644
--- a/lang/codegen/src/trait_definition.rs
+++ b/lang/codegen/src/trait_definition.rs
@@ -48,6 +48,9 @@ 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);
+    println!("mock_type: {:?}", mock_type);
+    println!("extracted_attrs: {:?}", attrs);
     let mut trait_item: ItemTrait = parse2(_input).unwrap();
     let trait_without_ink_attrs;
     let ink_code;
@@ -60,6 +63,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 +101,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 +117,12 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
                 #ink_trait
             }
         };
+
+        let pub_mock_env_ident = format_ident!("Mock{}", trait_item.ident.to_string());
+        maybe_use_mock_env = quote! {
+            #[cfg(test)]
+            pub use #namespace_ident :: mock_env as #pub_mock_env_ident;
+        };
     } else {
         trait_without_ink_attrs = trait_item;
         ink_code = quote! {};
@@ -124,7 +134,10 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
         #trait_without_ink_attrs
 
         #ink_code
+
+        #maybe_use_mock_env
     };
+    // println!("Final output: {}", code);
     code.into()
 }
 
@@ -167,7 +180,7 @@ 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![];
@@ -252,15 +265,33 @@ fn generate_wrapper(ink_trait: ItemTrait) -> proc_macro2::TokenStream {
                 >;
             });
 
+            let message_test_impl = match &mock_type {
+                Some(_mock_ty) => quote! {
+                    mock_env :: with(|mock_obj|
+                        mock_obj . #message_ident (
+                            #( #input_bindings , )*
+                        )
+                    ).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(test))]
+                    {
+                        Self::#message_builder_ident(self #( , #input_bindings)*)
+                            .fire()
+                            .unwrap_or_else(|err| ::core::panic!("{}: {:?}", #panic_str, err))
+                    }
+                    #[cfg(test)]
+                    {
+                        #message_test_impl
+                    }
                 }
 
                 #[inline]
@@ -292,6 +323,14 @@ 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(test)]
+            ::environmental::environmental!( pub mock_env : #ty );
+        },
+        None => quote! {},
+    };
+
     quote! {
         pub trait #trait_wrapper_ident {
             #( #def_messages )*
@@ -300,6 +339,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 +375,53 @@ fn remove_ink_attrs(mut trait_item: ItemTrait) -> ItemTrait {
     });
     trait_item
 }
+
+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)
+}
+
+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);
+    }
+}
\ No newline at end of file
diff --git a/lang/codegen/src/trait_definition/attr_args.rs b/lang/codegen/src/trait_definition/attr_args.rs
new file mode 100644
index 000000000..255b37033
--- /dev/null
+++ b/lang/codegen/src/trait_definition/attr_args.rs
@@ -0,0 +1,307 @@
+// Taken from ink_lang::ast::attr_args
+
+// Copyright 2018-2021 Parity Technologies (UK) Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use proc_macro2::{
+    Ident,
+    TokenStream as TokenStream2,
+};
+use quote::ToTokens;
+use syn::{
+    ext::IdentExt as _,
+    parse::{
+        Parse,
+        ParseStream,
+    },
+    punctuated::Punctuated,
+    spanned::Spanned,
+    Token,
+};
+
+/// The attribute arguments for the configuration of an ink! smart contract.
+///
+/// These are the segments `env = ::my::env::Environment` and `compile_as_dependency = true`
+/// in `#[ink::contract(env = ::my::env::Environment, compile_as_dependency = true`.
+#[derive(Debug, PartialEq, Eq)]
+pub struct AttributeArgs {
+    args: Punctuated<MetaNameValue, Token![,]>,
+}
+
+/// A name-value pair within an attribute, like `feature = "nightly"`.
+///
+/// The only difference from `syn::MetaNameValue` is that this additionally
+/// allows the `value` to be a plain identifier or path.
+#[derive(Debug, PartialEq, Eq)]
+pub struct MetaNameValue {
+    pub name: syn::Path,
+    pub eq_token: syn::token::Eq,
+    pub value: PathOrLit,
+}
+
+/// Either a path or a literal.
+#[derive(Debug, PartialEq, Eq)]
+pub enum PathOrLit {
+    Path(syn::Path),
+    Lit(syn::Lit),
+}
+
+impl IntoIterator for AttributeArgs {
+    type Item = MetaNameValue;
+    type IntoIter = syn::punctuated::IntoIter<MetaNameValue>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.args.into_iter()
+    }
+}
+
+impl Parse for AttributeArgs {
+    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
+        Ok(Self {
+            args: Punctuated::parse_terminated(input)?,
+        })
+    }
+}
+
+impl Parse for MetaNameValue {
+    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
+        let path = input.call(Self::parse_meta_path)?;
+        Self::parse_meta_name_value_after_path(path, input)
+    }
+}
+
+impl ToTokens for PathOrLit {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        match self {
+            Self::Lit(lit) => lit.to_tokens(tokens),
+            Self::Path(path) => path.to_tokens(tokens),
+        }
+    }
+}
+
+impl ToTokens for MetaNameValue {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        self.name.to_tokens(tokens);
+        self.eq_token.to_tokens(tokens);
+        self.value.to_tokens(tokens);
+    }
+}
+
+impl MetaNameValue {
+    /// Like [`syn::Path::parse_mod_style`] but accepts keywords in the path.
+    ///
+    /// # Note
+    ///
+    /// This code was taken from the `syn` implementation for a very similar
+    /// syntactical pattern.
+    fn parse_meta_path(input: ParseStream) -> Result<syn::Path, syn::Error> {
+        Ok(syn::Path {
+            leading_colon: input.parse()?,
+            segments: {
+                let mut segments = Punctuated::new();
+                while input.peek(Ident::peek_any) {
+                    let ident = Ident::parse_any(input)?;
+                    segments.push_value(syn::PathSegment::from(ident));
+                    if !input.peek(syn::Token![::]) {
+                        break
+                    }
+                    let punct = input.parse()?;
+                    segments.push_punct(punct);
+                }
+                if segments.is_empty() {
+                    return Err(input.error("expected path"))
+                } else if segments.trailing_punct() {
+                    return Err(input.error("expected path segment"))
+                }
+                segments
+            },
+        })
+    }
+
+    fn parse_meta_name_value_after_path(
+        name: syn::Path,
+        input: ParseStream,
+    ) -> Result<MetaNameValue, syn::Error> {
+        let span = name.span();
+        Ok(MetaNameValue {
+            name,
+            eq_token: input.parse().map_err(|_error| {
+                syn::Error::new(
+                    <_ as ::syn::spanned::Spanned>::span(&span),
+                    "ink! config options require an argument separated by '='",
+                )
+            })?,
+            value: input.parse()?,
+        })
+    }
+}
+
+impl Parse for PathOrLit {
+    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
+        if input.fork().peek(syn::Lit) {
+            return input.parse::<syn::Lit>().map(PathOrLit::Lit)
+        }
+        if input.fork().peek(Ident::peek_any) || input.fork().peek(Token![::]) {
+            return input.parse::<syn::Path>().map(PathOrLit::Path)
+        }
+        Err(input.error("cannot parse into either literal or path"))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use quote::quote;
+
+    impl AttributeArgs {
+        /// Creates a new attribute argument list from the given arguments.
+        pub fn new<I>(args: I) -> Self
+        where
+            I: IntoIterator<Item = MetaNameValue>,
+        {
+            Self {
+                args: args.into_iter().collect(),
+            }
+        }
+    }
+
+    #[test]
+    fn empty_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! {}).unwrap(),
+            AttributeArgs::new(vec![])
+        )
+    }
+
+    #[test]
+    fn literal_bool_value_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! { name = true }).unwrap(),
+            AttributeArgs::new(vec![MetaNameValue {
+                name: syn::parse_quote! { name },
+                eq_token: syn::parse_quote! { = },
+                value: PathOrLit::Lit(syn::parse_quote! { true }),
+            }])
+        )
+    }
+
+    #[test]
+    fn literal_str_value_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! { name = "string literal" }).unwrap(),
+            AttributeArgs::new(vec![MetaNameValue {
+                name: syn::parse_quote! { name },
+                eq_token: syn::parse_quote! { = },
+                value: PathOrLit::Lit(syn::parse_quote! { "string literal" }),
+            }])
+        )
+    }
+
+    #[test]
+    fn ident_value_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! { name = MyIdentifier }).unwrap(),
+            AttributeArgs::new(vec![MetaNameValue {
+                name: syn::parse_quote! { name },
+                eq_token: syn::parse_quote! { = },
+                value: PathOrLit::Path(syn::parse_quote! { MyIdentifier }),
+            }])
+        )
+    }
+
+    #[test]
+    fn root_path_value_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! { name = ::this::is::my::Path }).unwrap(),
+            AttributeArgs::new(vec![MetaNameValue {
+                name: syn::parse_quote! { name },
+                eq_token: syn::parse_quote! { = },
+                value: PathOrLit::Path(syn::parse_quote! { ::this::is::my::Path }),
+            }])
+        )
+    }
+
+    #[test]
+    fn relative_path_value_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! { name = this::is::my::relative::Path })
+                .unwrap(),
+            AttributeArgs::new(vec![MetaNameValue {
+                name: syn::parse_quote! { name },
+                eq_token: syn::parse_quote! { = },
+                value: PathOrLit::Path(
+                    syn::parse_quote! { this::is::my::relative::Path }
+                ),
+            }])
+        )
+    }
+
+    #[test]
+    fn trailing_comma_works() {
+        let mut expected_args = Punctuated::new();
+        expected_args.push_value(MetaNameValue {
+            name: syn::parse_quote! { name },
+            eq_token: syn::parse_quote! { = },
+            value: PathOrLit::Path(syn::parse_quote! { value }),
+        });
+        expected_args.push_punct(<Token![,]>::default());
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! { name = value, }).unwrap(),
+            AttributeArgs {
+                args: expected_args,
+            }
+        )
+    }
+
+    #[test]
+    fn many_mixed_works() {
+        assert_eq!(
+            syn::parse2::<AttributeArgs>(quote! {
+                name1 = ::root::Path,
+                name2 = false,
+                name3 = "string literal",
+                name4 = 42,
+                name5 = 7.7
+            })
+            .unwrap(),
+            AttributeArgs::new(vec![
+                MetaNameValue {
+                    name: syn::parse_quote! { name1 },
+                    eq_token: syn::parse_quote! { = },
+                    value: PathOrLit::Path(syn::parse_quote! { ::root::Path }),
+                },
+                MetaNameValue {
+                    name: syn::parse_quote! { name2 },
+                    eq_token: syn::parse_quote! { = },
+                    value: PathOrLit::Lit(syn::parse_quote! { false }),
+                },
+                MetaNameValue {
+                    name: syn::parse_quote! { name3 },
+                    eq_token: syn::parse_quote! { = },
+                    value: PathOrLit::Lit(syn::parse_quote! { "string literal" }),
+                },
+                MetaNameValue {
+                    name: syn::parse_quote! { name4 },
+                    eq_token: syn::parse_quote! { = },
+                    value: PathOrLit::Lit(syn::parse_quote! { 42 }),
+                },
+                MetaNameValue {
+                    name: syn::parse_quote! { name5 },
+                    eq_token: syn::parse_quote! { = },
+                    value: PathOrLit::Lit(syn::parse_quote! { 7.7 }),
+                },
+            ])
+        )
+    }
+}

From 51eaf214e7475c8aa3470fd6f7d5e675383027a0 Mon Sep 17 00:00:00 2001
From: h4x3rotab <h4x3rotab@gmail.com>
Date: Thu, 16 Jun 2022 08:11:17 +0000
Subject: [PATCH 2/5] macro to deploy mock contracts

---
 lang/codegen/src/trait_definition.rs | 60 ++++++++++++++++++++++++----
 1 file changed, 52 insertions(+), 8 deletions(-)

diff --git a/lang/codegen/src/trait_definition.rs b/lang/codegen/src/trait_definition.rs
index 2f39f2a7c..e8dfbe951 100644
--- a/lang/codegen/src/trait_definition.rs
+++ b/lang/codegen/src/trait_definition.rs
@@ -118,10 +118,12 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
             }
         };
 
-        let pub_mock_env_ident = format_ident!("Mock{}", trait_item.ident.to_string());
+        let pub_mock_env_ident = format_ident!("mock_{}", trait_item.ident.to_string().to_lowercase());
         maybe_use_mock_env = quote! {
             #[cfg(test)]
-            pub use #namespace_ident :: mock_env as #pub_mock_env_ident;
+            pub mod #pub_mock_env_ident {
+                pub use super :: #namespace_ident :: { mock_env as env , using , deploy };
+            }
         };
     } else {
         trait_without_ink_attrs = trait_item;
@@ -185,6 +187,7 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
     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()
@@ -267,11 +270,12 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
 
             let message_test_impl = match &mock_type {
                 Some(_mock_ty) => quote! {
-                    mock_env :: with(|mock_obj|
-                        mock_obj . #message_ident (
+                    mock_env :: with(|registry| {
+                        let mut mock_ref = registry.get_mut(self).expect("not an address of mocked contract");
+                        mock_ref.borrow_mut(). #message_ident (
                             #( #input_bindings , )*
                         )
-                    ).expect("mock object not set")
+                    }).expect("mock object not set")
                 },
                 None => quote! { ::core::panic!("cross-contract call is not supported in ink tests; try to set a mock object?") }
             };
@@ -324,9 +328,39 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
     let def_messages = def_messages.iter();
 
     let maybe_mock_environmental = match mock_type {
-        Some(ty) => quote! {
-            #[cfg(test)]
-            ::environmental::environmental!( pub mock_env : #ty );
+        Some(ty) =>{
+            quote! {
+                #[cfg(test)]
+                ::environmental::environmental!(
+                    pub mock_env : std::collections::BTreeMap<
+                        ::openbrush::traits::AccountId,
+                        std::rc::Rc<std::cell::RefCell< #ty >>
+                    > 
+                );
+
+                #[cfg(test)]
+                pub fn using<F: FnOnce()>(f: F) {
+                    let mut env = Default::default();
+                    mock_env::using(&mut env, f);
+                }
+
+                #[cfg(test)]
+                pub fn deploy(inner_contract : #ty) -> (::openbrush::traits::AccountId, std::rc::Rc<std::cell::RefCell< #ty >>) {
+                    let contract: std::rc::Rc<std::cell::RefCell< #ty >> = std::rc::Rc::new(
+                        std::cell::RefCell::< #ty >::new(inner_contract)
+                    );
+                    mock_env::with(|register| {
+                        let n: u8 = 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();
+
+                        register.insert(account_id.clone(), contract.clone());
+                        (account_id, contract)
+                    }).expect("must call within `using()`")
+                }
+            }
         },
         None => quote! {},
     };
@@ -376,6 +410,9 @@ 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");
 
@@ -395,6 +432,13 @@ fn extract_mock_config(attr: TokenStream) -> (Option<TokenStream>, TokenStream)
     (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)]

From e95b26f49f06cc87de61a20b0eee0b89ae8bef02 Mon Sep 17 00:00:00 2001
From: h4x3rotab <h4x3rotab@gmail.com>
Date: Tue, 21 Jun 2022 11:04:58 +0000
Subject: [PATCH 3/5] Addressable for auto stack management

---
 lang/codegen/src/trait_definition.rs |  89 +++++++-----
 lang/src/lib.rs                      |   2 +
 lang/src/traits.rs                   | 210 +++++++++++++++++++++++++++
 3 files changed, 266 insertions(+), 35 deletions(-)

diff --git a/lang/codegen/src/trait_definition.rs b/lang/codegen/src/trait_definition.rs
index e8dfbe951..4dc1078f0 100644
--- a/lang/codegen/src/trait_definition.rs
+++ b/lang/codegen/src/trait_definition.rs
@@ -120,7 +120,7 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
 
         let pub_mock_env_ident = format_ident!("mock_{}", trait_item.ident.to_string().to_lowercase());
         maybe_use_mock_env = quote! {
-            #[cfg(test)]
+            #[cfg(any(test, feature = "mockable"))]
             pub mod #pub_mock_env_ident {
                 pub use super :: #namespace_ident :: { mock_env as env , using , deploy };
             }
@@ -270,11 +270,14 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
 
             let message_test_impl = match &mock_type {
                 Some(_mock_ty) => quote! {
-                    mock_env :: with(|registry| {
-                        let mut mock_ref = registry.get_mut(self).expect("not an address of mocked contract");
-                        mock_ref.borrow_mut(). #message_ident (
+                    mock_env :: with(|ctx| {
+                        let mut mock_ref = ctx.register.get_mut(self).expect("not an address of mocked contract");
+                        ctx.stack.borrow_mut().push(&self);
+                        let result = mock_ref.borrow_mut(). #message_ident (
                             #( #input_bindings , )*
-                        )
+                        );
+                        ctx.stack.borrow_mut().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?") }
@@ -286,13 +289,13 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
                     & self
                     #( , #input_bindings : #input_types )*
                 ) -> #output_ty {
-                    #[cfg(not(test))]
+                    #[cfg(not(any(test, feature = "mockable")))]
                     {
                         Self::#message_builder_ident(self #( , #input_bindings)*)
                             .fire()
                             .unwrap_or_else(|err| ::core::panic!("{}: {:?}", #panic_str, err))
                     }
-                    #[cfg(test)]
+                    #[cfg(any(test, feature = "mockable"))]
                     {
                         #message_test_impl
                     }
@@ -328,40 +331,60 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
     let def_messages = def_messages.iter();
 
     let maybe_mock_environmental = match mock_type {
-        Some(ty) =>{
+        Some(ty) => {
             quote! {
-                #[cfg(test)]
-                ::environmental::environmental!(
-                    pub mock_env : std::collections::BTreeMap<
+                #[cfg(any(test, feature = "mockable"))]
+                pub struct Context {
+                    pub stack: std::rc::Rc<std::cell::RefCell<
+                        ::openbrush::traits::mock::ManagedCallStack
+                    >>,
+                    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(test)]
-                pub fn using<F: FnOnce()>(f: F) {
-                    let mut env = Default::default();
+                #[cfg(any(test, feature = "mockable"))]
+                pub fn using<F: FnOnce()>(
+                    stack: std::rc::Rc<std::cell::RefCell<::openbrush::traits::mock::ManagedCallStack>>,
+                    f: F
+                ) {
+                    let mut env = Context {
+                        stack,
+                        register: Default::default()
+                    };
                     mock_env::using(&mut env, f);
                 }
 
-                #[cfg(test)]
-                pub fn deploy(inner_contract : #ty) -> (::openbrush::traits::AccountId, std::rc::Rc<std::cell::RefCell< #ty >>) {
+                #[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)
                     );
-                    mock_env::with(|register| {
-                        let n: u8 = register.len().try_into()
+                    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();
 
-                        register.insert(account_id.clone(), contract.clone());
-                        (account_id, contract)
-                    }).expect("must call within `using()`")
+                        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! {},
     };
 
@@ -411,17 +434,13 @@ fn remove_ink_attrs(mut trait_item: ItemTrait) -> ItemTrait {
 }
 
 /// 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_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 }
@@ -452,20 +471,20 @@ mod tests {
                 mock = MyMockType,
                 namespace = ::name::space
             },
-            quote!{
+            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);
     }
-}
\ No newline at end of file
+}
diff --git a/lang/src/lib.rs b/lang/src/lib.rs
index a05d8d271..4e61893b2 100644
--- a/lang/src/lib.rs
+++ b/lang/src/lib.rs
@@ -21,6 +21,8 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
+extern crate alloc;
+
 pub mod derive;
 mod macros;
 pub mod storage;
diff --git a/lang/src/traits.rs b/lang/src/traits.rs
index 7aba40000..861ce2c09 100644
--- a/lang/src/traits.rs
+++ b/lang/src/traits.rs
@@ -77,3 +77,213 @@ pub trait Flush: ::ink_storage::traits::SpreadLayout + InkStorage {
 }
 
 impl<T: ::ink_storage::traits::SpreadLayout + InkStorage> Flush for T {}
+
+/// Types for managing mock cross-contract calls in unit tests
+pub mod mock {
+    use super::AccountId;
+
+    use alloc::{
+        rc::Rc,
+        vec::Vec,
+    };
+    use core::{
+        cell::{
+            Ref,
+            RefCell,
+            RefMut,
+        },
+        ops::{
+            Deref,
+            DerefMut,
+        },
+    };
+
+    /// A frame in the call stack
+    #[derive(Clone, Debug)]
+    pub struct MockCallContext {
+        pub level: u32,
+        pub caller: Option<AccountId>,
+        pub callee: AccountId,
+    }
+
+    /// A managed call stack for mocking cross-contract call in test environment
+    pub struct ManagedCallStack {
+        stack: Vec<MockCallContext>,
+    }
+
+    impl ManagedCallStack {
+        /// Crates a call stack with the default `account`
+        pub fn new(account: AccountId) -> Self {
+            ManagedCallStack {
+                stack: alloc::vec![MockCallContext {
+                    level: 0,
+                    caller: None,
+                    callee: account,
+                }],
+            }
+        }
+
+        /// Creates a call stack with the default `account` and returns a shared reference
+        pub fn create_shared(account: AccountId) -> Rc<RefCell<Self>> {
+            Rc::new(RefCell::new(Self::new(account)))
+        }
+
+        /// Changes the caller account
+        ///
+        /// Only allowed outside any contract call (when the stack is empty).
+        pub fn switch_account(&mut self, account: AccountId) -> Result<(), ()> {
+            if self.stack.len() != 1 {
+                return Err(())
+            }
+            let ctx = self.stack.get_mut(0).ok_or(())?;
+            ctx.callee = account;
+            Ok(())
+        }
+
+        /// Pushes a new call frame
+        pub fn push(&mut self, callee: &AccountId) {
+            let parent_ctx = self.peek().clone();
+            self.stack.push(MockCallContext {
+                level: parent_ctx.level + 1,
+                caller: Some(parent_ctx.callee),
+                callee: callee.clone(),
+            });
+            self.sync_to_ink();
+        }
+
+        /// Pops the call frame and returns the frame
+        pub fn pop(&mut self) -> Option<MockCallContext> {
+            if self.stack.len() > 1 {
+                let ctx = self.stack.pop();
+                self.sync_to_ink();
+                ctx
+            } else {
+                None
+            }
+        }
+
+        /// Peeks the current call frame
+        pub fn peek(&self) -> &MockCallContext {
+            self.stack.last().expect("stack is never empty; qed.")
+        }
+
+        /// Syncs the top call frame to ink testing environment
+        pub fn sync_to_ink(&self) {
+            let ctx = self.peek();
+            if let Some(caller) = ctx.caller {
+                ink_env::test::set_caller::<ink_env::DefaultEnvironment>(caller);
+            }
+            ink_env::test::set_callee::<ink_env::DefaultEnvironment>(ctx.callee);
+        }
+    }
+
+    /// A wrapper of a contract with an address for call stake auto-management
+    #[derive(Clone)]
+    pub struct Addressable<T> {
+        inner: Rc<RefCell<T>>,
+        id: AccountId,
+        stack: Rc<RefCell<ManagedCallStack>>,
+    }
+
+    impl<T> Addressable<T> {
+        /// Wraps a contract reference with id and a shared call stack
+        pub fn new(id: AccountId, inner: Rc<RefCell<T>>, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            Addressable { inner, id, stack }
+        }
+
+        /// Wraps a native contract object with a simple id
+        ///
+        /// The account id of the contract will be the `id` with zero-padding.
+        pub fn create_native(id: u8, inner: T, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            Addressable {
+                inner: Rc::new(RefCell::new(inner)),
+                id: naive_id(id),
+                stack,
+            }
+        }
+
+        /// Returns the account id of the inner contract
+        pub fn id(&self) -> AccountId {
+            self.id.clone()
+        }
+
+        /// Borrows the contract for _a_ call with the stack auto-managed
+        ///
+        /// Holding the ref for multiple calls or nested call is considered abuse.
+        pub fn call(&self) -> ScopedRef<'_, T> {
+            ScopedRef::new(self.inner.borrow(), &self.id, self.stack.clone())
+        }
+
+        /// Borrows the contract for _a_ mut call with the stack auto-managed
+        ///
+        /// Holding the mut ref for multiple calls or nested call is considered abuse.
+        pub fn call_mut(&self) -> ScopedRefMut<'_, T> {
+            ScopedRefMut::new(self.inner.borrow_mut(), &self.id, self.stack.clone())
+        }
+    }
+
+    /// Push a call stack when the `Ref` in scope
+    pub struct ScopedRef<'b, T: 'b> {
+        inner: Ref<'b, T>,
+        stack: Rc<RefCell<ManagedCallStack>>,
+    }
+
+    impl<'b, T> ScopedRef<'b, T> {
+        fn new(inner: Ref<'b, T>, address: &AccountId, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            stack.borrow_mut().push(address);
+            Self { inner, stack }
+        }
+    }
+
+    impl<'b, T> Deref for ScopedRef<'b, T> {
+        type Target = T;
+        fn deref(&self) -> &T {
+            self.inner.deref()
+        }
+    }
+
+    impl<'b, T> Drop for ScopedRef<'b, T> {
+        fn drop(&mut self) {
+            self.stack.borrow_mut().pop().expect("pop never fails");
+        }
+    }
+
+    /// Push a call stack when the `RefMut` in scope
+    pub struct ScopedRefMut<'b, T: 'b> {
+        inner: RefMut<'b, T>,
+        stack: Rc<RefCell<ManagedCallStack>>,
+    }
+
+    impl<'b, T> ScopedRefMut<'b, T> {
+        fn new(inner: RefMut<'b, T>, address: &AccountId, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            stack.borrow_mut().push(address);
+            Self { inner, stack }
+        }
+    }
+
+    impl<'b, T> Deref for ScopedRefMut<'b, T> {
+        type Target = T;
+        fn deref(&self) -> &T {
+            self.inner.deref()
+        }
+    }
+
+    impl<'b, T> DerefMut for ScopedRefMut<'b, T> {
+        fn deref_mut(&mut self) -> &mut T {
+            self.inner.deref_mut()
+        }
+    }
+
+    impl<'b, T> Drop for ScopedRefMut<'b, T> {
+        fn drop(&mut self) {
+            self.stack.borrow_mut().pop().expect("pop never fails");
+        }
+    }
+
+    /// Generates a naive zero-padding account id with a `u8` number
+    pub fn naive_id(id: u8) -> AccountId {
+        let mut address = [0u8; 32];
+        address[31] = id;
+        address.into()
+    }
+}

From 5a8ac927eb40d521bd1414fac72cff0ef43b5e3b Mon Sep 17 00:00:00 2001
From: h4x3rotab <h4x3rotab@gmail.com>
Date: Wed, 22 Jun 2022 05:53:10 +0000
Subject: [PATCH 4/5] make ManagedCallStack a shared prt

---
 lang/codegen/src/trait_definition.rs | 10 ++---
 lang/src/traits.rs                   | 63 +++++++++++++---------------
 2 files changed, 34 insertions(+), 39 deletions(-)

diff --git a/lang/codegen/src/trait_definition.rs b/lang/codegen/src/trait_definition.rs
index 4dc1078f0..ca4735437 100644
--- a/lang/codegen/src/trait_definition.rs
+++ b/lang/codegen/src/trait_definition.rs
@@ -272,11 +272,11 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
                 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.borrow_mut().push(&self);
+                        ctx.stack.push(&self);
                         let result = mock_ref.borrow_mut(). #message_ident (
                             #( #input_bindings , )*
                         );
-                        ctx.stack.borrow_mut().pop();
+                        ctx.stack.pop();
                         result
                     }).expect("mock object not set")
                 },
@@ -335,9 +335,7 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
             quote! {
                 #[cfg(any(test, feature = "mockable"))]
                 pub struct Context {
-                    pub stack: std::rc::Rc<std::cell::RefCell<
-                        ::openbrush::traits::mock::ManagedCallStack
-                    >>,
+                    pub stack: ::openbrush::traits::mock::SharedCallStack,
                     pub register: std::collections::BTreeMap<
                         ::openbrush::traits::AccountId,
                         std::rc::Rc<std::cell::RefCell< #ty >>
@@ -351,7 +349,7 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
 
                 #[cfg(any(test, feature = "mockable"))]
                 pub fn using<F: FnOnce()>(
-                    stack: std::rc::Rc<std::cell::RefCell<::openbrush::traits::mock::ManagedCallStack>>,
+                    stack: ::openbrush::traits::mock::SharedCallStack,
                     f: F
                 ) {
                     let mut env = Context {
diff --git a/lang/src/traits.rs b/lang/src/traits.rs
index 861ce2c09..443d2bf62 100644
--- a/lang/src/traits.rs
+++ b/lang/src/traits.rs
@@ -107,43 +107,40 @@ pub mod mock {
     }
 
     /// A managed call stack for mocking cross-contract call in test environment
-    pub struct ManagedCallStack {
-        stack: Vec<MockCallContext>,
+    #[derive(Clone)]
+    pub struct SharedCallStack {
+        stack: Rc<RefCell<Vec<MockCallContext>>>,
     }
 
-    impl ManagedCallStack {
+    impl SharedCallStack {
         /// Crates a call stack with the default `account`
         pub fn new(account: AccountId) -> Self {
-            ManagedCallStack {
-                stack: alloc::vec![MockCallContext {
+            SharedCallStack {
+                stack: Rc::new(RefCell::new(alloc::vec![MockCallContext {
                     level: 0,
                     caller: None,
                     callee: account,
-                }],
+                }])),
             }
         }
 
-        /// Creates a call stack with the default `account` and returns a shared reference
-        pub fn create_shared(account: AccountId) -> Rc<RefCell<Self>> {
-            Rc::new(RefCell::new(Self::new(account)))
-        }
-
         /// Changes the caller account
         ///
         /// Only allowed outside any contract call (when the stack is empty).
-        pub fn switch_account(&mut self, account: AccountId) -> Result<(), ()> {
-            if self.stack.len() != 1 {
+        pub fn switch_account(&self, account: AccountId) -> Result<(), ()> {
+            let mut stack = self.stack.borrow_mut();
+            if stack.len() != 1 {
                 return Err(())
             }
-            let ctx = self.stack.get_mut(0).ok_or(())?;
+            let ctx = stack.get_mut(0).ok_or(())?;
             ctx.callee = account;
             Ok(())
         }
 
         /// Pushes a new call frame
-        pub fn push(&mut self, callee: &AccountId) {
-            let parent_ctx = self.peek().clone();
-            self.stack.push(MockCallContext {
+        pub fn push(&self, callee: &AccountId) {
+            let parent_ctx = self.peek();
+            self.stack.borrow_mut().push(MockCallContext {
                 level: parent_ctx.level + 1,
                 caller: Some(parent_ctx.callee),
                 callee: callee.clone(),
@@ -152,9 +149,9 @@ pub mod mock {
         }
 
         /// Pops the call frame and returns the frame
-        pub fn pop(&mut self) -> Option<MockCallContext> {
-            if self.stack.len() > 1 {
-                let ctx = self.stack.pop();
+        pub fn pop(&self) -> Option<MockCallContext> {
+            if self.stack.borrow().len() > 1 {
+                let ctx = self.stack.borrow_mut().pop();
                 self.sync_to_ink();
                 ctx
             } else {
@@ -163,8 +160,8 @@ pub mod mock {
         }
 
         /// Peeks the current call frame
-        pub fn peek(&self) -> &MockCallContext {
-            self.stack.last().expect("stack is never empty; qed.")
+        pub fn peek(&self) -> MockCallContext {
+            self.stack.borrow().last().cloned().expect("stack is never empty; qed.")
         }
 
         /// Syncs the top call frame to ink testing environment
@@ -182,19 +179,19 @@ pub mod mock {
     pub struct Addressable<T> {
         inner: Rc<RefCell<T>>,
         id: AccountId,
-        stack: Rc<RefCell<ManagedCallStack>>,
+        stack: SharedCallStack,
     }
 
     impl<T> Addressable<T> {
         /// Wraps a contract reference with id and a shared call stack
-        pub fn new(id: AccountId, inner: Rc<RefCell<T>>, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+        pub fn new(id: AccountId, inner: Rc<RefCell<T>>, stack: SharedCallStack) -> Self {
             Addressable { inner, id, stack }
         }
 
         /// Wraps a native contract object with a simple id
         ///
         /// The account id of the contract will be the `id` with zero-padding.
-        pub fn create_native(id: u8, inner: T, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+        pub fn create_native(id: u8, inner: T, stack: SharedCallStack) -> Self {
             Addressable {
                 inner: Rc::new(RefCell::new(inner)),
                 id: naive_id(id),
@@ -225,12 +222,12 @@ pub mod mock {
     /// Push a call stack when the `Ref` in scope
     pub struct ScopedRef<'b, T: 'b> {
         inner: Ref<'b, T>,
-        stack: Rc<RefCell<ManagedCallStack>>,
+        stack: SharedCallStack,
     }
 
     impl<'b, T> ScopedRef<'b, T> {
-        fn new(inner: Ref<'b, T>, address: &AccountId, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
-            stack.borrow_mut().push(address);
+        fn new(inner: Ref<'b, T>, address: &AccountId, stack: SharedCallStack) -> Self {
+            stack.push(address);
             Self { inner, stack }
         }
     }
@@ -244,19 +241,19 @@ pub mod mock {
 
     impl<'b, T> Drop for ScopedRef<'b, T> {
         fn drop(&mut self) {
-            self.stack.borrow_mut().pop().expect("pop never fails");
+            self.stack.pop().expect("pop never fails");
         }
     }
 
     /// Push a call stack when the `RefMut` in scope
     pub struct ScopedRefMut<'b, T: 'b> {
         inner: RefMut<'b, T>,
-        stack: Rc<RefCell<ManagedCallStack>>,
+        stack: SharedCallStack,
     }
 
     impl<'b, T> ScopedRefMut<'b, T> {
-        fn new(inner: RefMut<'b, T>, address: &AccountId, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
-            stack.borrow_mut().push(address);
+        fn new(inner: RefMut<'b, T>, address: &AccountId, stack: SharedCallStack) -> Self {
+            stack.push(address);
             Self { inner, stack }
         }
     }
@@ -276,7 +273,7 @@ pub mod mock {
 
     impl<'b, T> Drop for ScopedRefMut<'b, T> {
         fn drop(&mut self) {
-            self.stack.borrow_mut().pop().expect("pop never fails");
+            self.stack.pop().expect("pop never fails");
         }
     }
 

From 69c2f5e9b7f719987c43b7b1778d444c0bf9cbde Mon Sep 17 00:00:00 2001
From: h4x3rotab <h4x3rotab@gmail.com>
Date: Wed, 22 Jun 2022 11:33:37 +0000
Subject: [PATCH 5/5] short circuit test-only code by mockable feature

---
 Cargo.toml                           | 1 +
 lang/Cargo.toml                      | 3 ++-
 lang/codegen/src/trait_definition.rs | 3 ---
 lang/src/traits.rs                   | 1 +
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6eca9e7d8..5f78a95b7 100644
--- a/Cargo.toml
+++ b/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"
diff --git a/lang/Cargo.toml b/lang/Cargo.toml
index beecfc946..5a818f7dd 100644
--- a/lang/Cargo.toml
+++ b/lang/Cargo.toml
@@ -43,4 +43,5 @@ std = [
     "ink_metadata/std",
     "scale/std",
     "scale-info/std",
-]
\ No newline at end of file
+]
+mockable = []
diff --git a/lang/codegen/src/trait_definition.rs b/lang/codegen/src/trait_definition.rs
index ca4735437..232ae984d 100644
--- a/lang/codegen/src/trait_definition.rs
+++ b/lang/codegen/src/trait_definition.rs
@@ -49,8 +49,6 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
     }
     let attrs: proc_macro2::TokenStream = _attrs.into();
     let (mock_type, attrs) = extract_mock_config(attrs);
-    println!("mock_type: {:?}", mock_type);
-    println!("extracted_attrs: {:?}", attrs);
     let mut trait_item: ItemTrait = parse2(_input).unwrap();
     let trait_without_ink_attrs;
     let ink_code;
@@ -139,7 +137,6 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
 
         #maybe_use_mock_env
     };
-    // println!("Final output: {}", code);
     code.into()
 }
 
diff --git a/lang/src/traits.rs b/lang/src/traits.rs
index 443d2bf62..aa7f14c4e 100644
--- a/lang/src/traits.rs
+++ b/lang/src/traits.rs
@@ -79,6 +79,7 @@ pub trait Flush: ::ink_storage::traits::SpreadLayout + InkStorage {
 impl<T: ::ink_storage::traits::SpreadLayout + InkStorage> Flush for T {}
 
 /// Types for managing mock cross-contract calls in unit tests
+#[cfg(feature = "mockable")]
 pub mod mock {
     use super::AccountId;