Update SandboxConfig

Cargo.toml

@@ -21,6 +21,7 @@ repository = "https://github.com/Cardinal-Cryptography/drink"
 version = "0.10.0"
 
 [workspace.dependencies]
+paste = { version = "1.0.7" }
 anyhow = { version = "1.0.71" }
 cargo_metadata = { version = "0.18.1" }
 clap = { version = "4.3.4" }

drink-cli/src/executor/mod.rs

@@ -80,7 +80,7 @@ fn add_tokens(app_state: &mut AppState, recipient: AccountId32, value: u128) ->
     app_state
         .session
         .sandbox()
-        .mint_into(recipient.clone(), value)
+        .mint_into(&recipient, value)
         .map_err(|err| anyhow::format_err!("Failed to add token: {err:?}"))?;
     app_state.print(&format!("{value} tokens added to {recipient}",));
     Ok(())

drink/Cargo.toml

@@ -10,6 +10,7 @@ version.workspace = true
 description = "Minimal sufficient architecture that allows for a fully functional ink! contract development"
 
 [dependencies]
+paste = { workspace = true }
 contract-metadata = { workspace = true, optional = true}
 contract-transcode = { workspace = true, optional = true }
 frame-metadata = { workspace = true }

drink/src/lib.rs

@@ -26,8 +26,11 @@ pub use mock::{mock_message, ContractMock, MessageMock, MockedCallResult, Select
 use pallet_contracts::{debug::ExecResult, ExecReturnValue};
 use pallet_contracts_uapi::ReturnFlags;
 use parity_scale_codec::{Decode, Encode};
-/// Export pallets that are used in the minimal runtime.
-pub use {frame_support, frame_system, pallet_balances, pallet_contracts, pallet_timestamp};
+/// Export crates that are used in the create_minimal_runtime macro.
+pub use {
+    frame_support, frame_system, pallet_balances, pallet_contracts, pallet_timestamp, paste,
+    sp_externalities::Extension, sp_io::TestExternalities,
+};
 
 pub use crate::runtime::minimal::{self, MinimalRuntime};
 use crate::{

drink/src/runtime/minimal.rs

@@ -1,65 +1,113 @@
 #![allow(missing_docs)] // `construct_macro` doesn't allow doc comments for the runtime type.
 
+use std::time::SystemTime;
+
+use frame_support::{sp_runtime::traits::Header, traits::Hooks};
+pub struct BlockBuilder<T>(std::marker::PhantomData<T>);
+
+impl<
+        T: pallet_balances::Config + pallet_timestamp::Config<Moment = u64> + pallet_contracts::Config,
+    > BlockBuilder<T>
+{
+    pub fn initialize_block(
+        height: frame_system::pallet_prelude::BlockNumberFor<T>,
+        parent_hash: <T as frame_system::Config>::Hash,
+    ) {
+        frame_system::Pallet::<T>::reset_events();
+        frame_system::Pallet::<T>::initialize(&height, &parent_hash, &Default::default());
+        pallet_balances::Pallet::<T>::on_initialize(height);
+        pallet_timestamp::Pallet::<T>::set_timestamp(
+            SystemTime::now()
+                .duration_since(SystemTime::UNIX_EPOCH)
+                .expect("Time went backwards")
+                .as_secs(),
+        );
+        pallet_timestamp::Pallet::<T>::on_initialize(height);
+        pallet_contracts::Pallet::<T>::on_initialize(height);
+        frame_system::Pallet::<T>::note_finished_initialize();
+    }
+
+    pub fn finalize_block(
+        height: frame_system::pallet_prelude::BlockNumberFor<T>,
+    ) -> Result<<T as frame_system::Config>::Hash, String> {
+        pallet_contracts::Pallet::<T>::on_finalize(height);
+        pallet_timestamp::Pallet::<T>::on_finalize(height);
+        pallet_balances::Pallet::<T>::on_finalize(height);
+        Ok(frame_system::Pallet::<T>::finalize().hash())
+    }
+}
+
 /// The macro will generate an implementation of `drink::SandboxConfig` for the given runtime type.
 #[macro_export]
 macro_rules! impl_sandbox_config {
-    (
-        $( #[ $attr:meta ] )*
-        $vis:vis struct $name:ident {
-            runtime: $runtime:tt;
-            default_balance: $default_balance:expr;
-            default_actor: $default_actor:expr;
+    ($name:ident, $runtime:tt, $default_balance:expr, $default_actor:expr) => {
+        $crate::paste::paste! {
+            $crate::impl_sandbox_config!([<EXT_ $name:upper>], $name, $runtime, $default_balance, $default_actor);
         }
-    ) => {
-        $( #[ $attr ] )*
-        $vis struct $name;
-        impl_sandbox_config!($name, $runtime, $default_balance, $default_actor);
     };
-    (
-        $name:ident, $runtime:tt, $default_balance:expr, $default_actor:expr
-    ) => {
-        impl $crate::SandboxConfig for $name {
-            type Runtime = $runtime;
 
-            fn initialize_storage(storage: &mut $crate::frame_support::sp_runtime::Storage) -> Result<(), String> {
+    ($ext:ident, $name:ident, $runtime:tt, $default_balance:expr, $default_actor:expr) => {
+        thread_local! {
+            static $ext: ::std::cell::RefCell<$crate::TestExternalities> = ::std::cell::RefCell::new({
                 use $crate::frame_support::sp_runtime::BuildStorage;
+                let mut storage = $crate::frame_system::GenesisConfig::<$runtime>::default()
+                    .build_storage().unwrap();
+
                 $crate::pallet_balances::GenesisConfig::<$runtime> {
-                    balances: vec![(Self::default_actor(), $default_balance)],
+                    balances: vec![($default_actor, $default_balance)],
                 }
-                .assimilate_storage(storage)
+                .assimilate_storage(&mut storage)
+                .unwrap();
+
+                let mut ext = $crate::TestExternalities::new(storage);
+                ext.execute_with(|| {
+                    <$name as $crate::SandboxConfig>::initialize_block(1, Default::default())
+                });
+
+                ext
+            });
+        }
+
+        impl $crate::SandboxConfig for $name {
+            type Runtime = $runtime;
+
+            fn execute_with<T>(execute: impl FnOnce() -> T) -> T {
+                $ext.with(|v| v.borrow_mut().execute_with(execute))
+            }
+
+            fn dry_run<T>(action: impl FnOnce() -> T) -> T {
+                $ext.with(|v| {
+                    // Make a backup of the backend.
+                    let backend_backup = v.borrow_mut().as_backend();
+
+                    // Run the action, potentially modifying storage. Ensure, that there are no pending changes
+                    // that would affect the reverted backend.
+                    let result = action();
+
+                    let mut v = v.borrow_mut();
+                    v.commit_all().expect("Failed to commit changes");
+
+                    // Restore the backend.
+                    v.backend = backend_backup;
+                    result
+                })
+            }
+
+            fn register_extension<E: ::core::any::Any + $crate::Extension>(ext: E) {
+                $ext.with(|v| v.borrow_mut().register_extension(ext));
             }
 
             fn initialize_block(
                 height: $crate::frame_system::pallet_prelude::BlockNumberFor<$runtime>,
                 parent_hash: <$runtime as $crate::frame_system::Config>::Hash,
-            ) -> Result<(), String> {
-                use std::time::SystemTime;
-                use $crate::frame_support::traits::Hooks;
-
-                $crate::frame_system::Pallet::<$runtime>::reset_events();
-                $crate::frame_system::Pallet::<$runtime>::initialize(&height, &parent_hash, &Default::default());
-                $crate::pallet_balances::Pallet::<$runtime>::on_initialize(height);
-                $crate::pallet_timestamp::Pallet::<$runtime>::set_timestamp(
-                    SystemTime::now()
-                        .duration_since(SystemTime::UNIX_EPOCH)
-                        .expect("Time went backwards")
-                        .as_secs(),
-                );
-                $crate::pallet_timestamp::Pallet::<$runtime>::on_initialize(height);
-                $crate::pallet_contracts::Pallet::<$runtime>::on_initialize(height);
-                $crate::frame_system::Pallet::<$runtime>::note_finished_initialize();
-                Ok(())
+            ) {
+                $crate::minimal::BlockBuilder::<$runtime>::initialize_block(height, parent_hash)
             }
 
             fn finalize_block(
                 height: $crate::frame_system::pallet_prelude::BlockNumberFor<$runtime>,
             ) -> Result<<$runtime as $crate::frame_system::Config>::Hash, String> {
-                use $crate::frame_support::traits::Hooks;
-
-                $crate::pallet_contracts::Pallet::<$runtime>::on_finalize(height);
-                $crate::pallet_timestamp::Pallet::<$runtime>::on_finalize(height);
-                $crate::pallet_balances::Pallet::<$runtime>::on_finalize(height);
-                Ok($crate::frame_system::Pallet::<$runtime>::finalize().hash())
+                $crate::minimal::BlockBuilder::<$runtime>::finalize_block(height)
             }
 
             fn default_actor() -> $crate::runtime::AccountIdFor<$runtime> {
@@ -211,6 +259,8 @@ mod construct_runtime {
 
     /// Default initial balance for the default account.
     pub const INITIAL_BALANCE: u128 = 1_000_000_000_000_000;
+
+    // implement the `drink::SandboxConfig` for the runtime type
     $crate::impl_sandbox_config!($name, $name, INITIAL_BALANCE, AccountId32::new([1u8; 32]));
 }
 

drink/src/sandbox.rs

@@ -11,46 +11,56 @@ pub mod timestamp_api;
 use std::any::Any;
 
 use sp_externalities::Extension;
-use sp_io::TestExternalities;
 
 /// A sandboxed runtime.
-pub struct Sandbox<Config> {
-    externalities: TestExternalities,
-    _phantom: std::marker::PhantomData<Config>,
-}
+#[derive(frame_support::DefaultNoBound)]
+pub struct Sandbox<Config>(std::marker::PhantomData<Config>);
 
-impl<Config> Sandbox<Config> {
+impl<Config: SandboxConfig> Sandbox<Config> {
     /// Execute the given closure with the inner externallities.
     ///
     /// Returns the result of the given closure.
-    pub fn execute_with<T>(&mut self, execute: impl FnOnce() -> T) -> T {
-        self.externalities.execute_with(execute)
+    pub fn execute_with<T>(&self, execute: impl FnOnce() -> T) -> T {
+        Config::execute_with(execute)
     }
 
     /// Run an action without modifying the storage.
     ///
     /// # Arguments
     ///
     /// * `action` - The action to run.
-    pub fn dry_run<T>(&mut self, action: impl FnOnce(&mut Self) -> T) -> T {
-        // Make a backup of the backend.
-        let backend_backup = self.externalities.as_backend();
+    pub fn dry_run<T>(&self, action: impl FnOnce() -> T) -> T {
+        Config::dry_run(action)
+    }
 
-        // Run the action, potentially modifying storage. Ensure, that there are no pending changes
-        // that would affect the reverted backend.
-        let result = action(self);
-        self.externalities
-            .commit_all()
-            .expect("Failed to commit changes");
+    /// Registers an extension.
+    pub fn register_extension<E: Any + Extension>(&self, ext: E) {
+        Config::register_extension(ext);
+    }
+}
 
-        // Restore the backend.
-        self.externalities.backend = backend_backup;
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::MinimalRuntime;
+    #[test]
+    fn dry_run_works() {
+        let sandbox = Sandbox::<MinimalRuntime>::default();
+        let balance = sandbox.free_balance(&MinimalRuntime::default_actor());
 
-        result
-    }
+        let dry_run_balance = sandbox.dry_run(|| {
+            sandbox
+                .mint_into(&MinimalRuntime::default_actor(), 100)
+                .unwrap();
 
-    /// Registers an extension.
-    pub fn register_extension<E: Any + Extension>(&mut self, ext: E) {
-        self.externalities.register_extension(ext);
+            sandbox.free_balance(&MinimalRuntime::default_actor())
+        });
+
+        assert_eq!(balance + 100, dry_run_balance);
+
+        assert_eq!(
+            sandbox.free_balance(&MinimalRuntime::default_actor()),
+            balance
+        );
     }
 }

drink/src/sandbox/balance_api.rs

@@ -15,13 +15,11 @@ where
     /// * `address` - The address of the account to add tokens to.
     /// * `amount` - The number of tokens to add.
     pub fn mint_into(
-        &mut self,
-        address: AccountIdFor<Config::Runtime>,
+        &self,
+        address: &AccountIdFor<Config::Runtime>,
         amount: BalanceOf<Config::Runtime>,
     ) -> Result<BalanceOf<Config::Runtime>, DispatchError> {
-        self.execute_with(|| {
-            pallet_balances::Pallet::<Config::Runtime>::mint_into(&address, amount)
-        })
+        self.execute_with(|| pallet_balances::Pallet::<Config::Runtime>::mint_into(address, amount))
     }
 
     /// Return the free balance of an account.
@@ -30,9 +28,29 @@ where
     ///
     /// * `address` - The address of the account to query.
     pub fn free_balance(
-        &mut self,
+        &self,
         address: &AccountIdFor<Config::Runtime>,
     ) -> BalanceOf<Config::Runtime> {
         self.execute_with(|| pallet_balances::Pallet::<Config::Runtime>::free_balance(address))
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::MinimalRuntime;
+    #[test]
+    fn mint_works() {
+        let sandbox = Sandbox::<MinimalRuntime>::default();
+        let balance = sandbox.free_balance(&MinimalRuntime::default_actor());
+
+        sandbox
+            .mint_into(&MinimalRuntime::default_actor(), 100)
+            .unwrap();
+
+        assert_eq!(
+            sandbox.free_balance(&MinimalRuntime::default_actor()),
+            balance + 100
+        );
+    }
+}