Basic contract mocking

.github/workflows/rust-checks.yml

@@ -59,6 +59,8 @@ jobs:
       - name: Run tests for examples
         shell: bash
         run: |
+          # todo: use loop xD
+
           pushd examples/flipper
           cargo contract build --release
           cargo test --release
@@ -73,3 +75,8 @@ jobs:
           cargo contract build --release
           cargo test --release
           popd
+
+          pushd examples/mocking
+          cargo contract build --release
+          cargo test --release
+          popd

Cargo.toml

@@ -10,6 +10,7 @@ exclude = [
     "examples/counter",
     "examples/flipper",
     "examples/cross-contract-call-tracing",
+    "examples/mocking",
 ]
 
 [workspace.package]
@@ -19,7 +20,7 @@ homepage = "https://github.com/Cardinal-Cryptography/drink"
 license = "Apache-2.0"
 readme = "README.md"
 repository = "https://github.com/Cardinal-Cryptography/drink"
-version = "0.4.1"
+version = "0.5.0"
 
 [workspace.dependencies]
 anyhow = { version = "1.0.71" }
@@ -40,7 +41,7 @@ frame-metadata = { version = "16.0.0" }
 frame-support = { version = "23.0.0" }
 frame-system = { version = "23.0.0" }
 pallet-balances = { version = "23.0.0" }
-pallet-contracts = { package = "pallet-contracts-for-drink", version = "22.0.0" }
+pallet-contracts = { package = "pallet-contracts-for-drink", version = "22.0.1" }
 pallet-contracts-primitives = { version = "26.0.0" }
 pallet-timestamp = { version = "22.0.0" }
 sp-core = { version = "23.0.0" }
@@ -50,4 +51,4 @@ sp-runtime-interface = { version = "19.0.0" }
 
 # Local dependencies
 
-drink = { version = "0.4.1", path = "drink" }
+drink = { version = "0.5.0", path = "drink" }

drink/Cargo.toml

@@ -26,8 +26,6 @@ sp-runtime-interface = { workspace = true }
 
 scale-info = { workspace = true }
 thiserror = { workspace = true }
-
-[dev-dependencies]
 wat = { workspace = true }
 
 [features]

drink/src/errors.rs

@@ -0,0 +1,43 @@
+//! Module gathering common error and result types.
+
+use thiserror::Error;
+
+/// Main error type for the drink crate.
+#[derive(Error, Debug)]
+pub enum Error {
+    /// Externalities could not be initialized.
+    #[error("Failed to build storage: {0}")]
+    StorageBuilding(String),
+    /// Block couldn't have been initialized.
+    #[error("Failed to initialize block: {0}")]
+    BlockInitialize(String),
+    /// Block couldn't have been finalized.
+    #[error("Failed to finalize block: {0}")]
+    BlockFinalize(String),
+}
+
+/// Every contract message wraps its return value in `Result<T, LangResult>`. This is the error
+/// type.
+///
+/// Copied from ink primitives.
+#[non_exhaustive]
+#[repr(u32)]
+#[derive(
+    Debug,
+    Copy,
+    Clone,
+    PartialEq,
+    Eq,
+    parity_scale_codec::Encode,
+    parity_scale_codec::Decode,
+    scale_info::TypeInfo,
+    Error,
+)]
+pub enum LangError {
+    /// Failed to read execution input for the dispatchable.
+    #[error("Failed to read execution input for the dispatchable.")]
+    CouldNotReadInput = 1u32,
+}
+
+/// The `Result` type for ink! messages.
+pub type MessageResult<T> = Result<T, LangError>;

drink/src/lib.rs

@@ -5,24 +5,37 @@
 
 pub mod chain_api;
 pub mod contract_api;
-mod error;
+pub mod errors;
+mod mock;
 pub mod runtime;
 #[cfg(feature = "session")]
 pub mod session;
-use std::marker::PhantomData;
 
-pub use error::Error;
+use std::{
+    marker::PhantomData,
+    sync::{Arc, Mutex},
+};
+
+pub use errors::Error;
 use frame_support::sp_runtime::{traits::One, BuildStorage};
 pub use frame_support::{
     sp_runtime::{AccountId32, DispatchError},
     weights::Weight,
 };
 use frame_system::{pallet_prelude::BlockNumberFor, EventRecord, GenesisConfig};
+pub use mock::{mock_message, ContractMock, MessageMock, MockedCallResult, MockingApi, Selector};
+use pallet_contracts::debug::ExecResult;
+use pallet_contracts_primitives::{ExecReturnValue, ReturnFlags};
+use parity_scale_codec::{Decode, Encode};
 use sp_io::TestExternalities;
 
 use crate::{
-    pallet_contracts_debugging::DebugExt,
-    runtime::{pallet_contracts_debugging::NoopDebugExt, *},
+    mock::MockRegistry,
+    pallet_contracts_debugging::{InterceptingExt, TracingExt},
+    runtime::{
+        pallet_contracts_debugging::{InterceptingExtT, NoopExt},
+        *,
+    },
 };
 
 /// Main result type for the drink crate.
@@ -35,6 +48,8 @@ pub type EventRecordOf<T> =
 /// A sandboxed runtime.
 pub struct Sandbox<R: Runtime> {
     externalities: TestExternalities,
+    mock_registry: Arc<Mutex<MockRegistry<AccountIdFor<R>>>>,
+    mock_counter: usize,
     _phantom: PhantomData<R>,
 }
 
@@ -57,6 +72,8 @@ impl<R: Runtime> Sandbox<R> {
 
         let mut sandbox = Self {
             externalities: TestExternalities::new(storage),
+            mock_registry: Arc::new(Mutex::new(MockRegistry::new())),
+            mock_counter: 0,
             _phantom: PhantomData,
         };
 
@@ -68,7 +85,9 @@ impl<R: Runtime> Sandbox<R> {
             .map_err(Error::BlockInitialize)?;
 
         // We register a noop debug extension by default.
-        sandbox.override_debug_handle(DebugExt(Box::new(NoopDebugExt {})));
+        sandbox.override_debug_handle(TracingExt(Box::new(NoopExt {})));
+
+        sandbox.setup_mock_extension();
 
         Ok(sandbox)
     }
@@ -77,7 +96,65 @@ impl<R: Runtime> Sandbox<R> {
     ///
     /// By default, a new `Sandbox` instance is created with a noop debug extension. This method
     /// allows to override it with a custom debug extension.
-    pub fn override_debug_handle(&mut self, d: DebugExt) {
+    pub fn override_debug_handle(&mut self, d: TracingExt) {
         self.externalities.register_extension(d);
     }
+
+    /// Registers the extension for intercepting calls to contracts.
+    fn setup_mock_extension(&mut self) {
+        self.externalities
+            .register_extension(InterceptingExt(Box::new(MockingExtension {
+                mock_registry: Arc::clone(&self.mock_registry),
+            })));
+    }
+}
+
+/// Runtime extension enabling contract call interception.
+struct MockingExtension<AccountId: Ord> {
+    /// Mock registry, shared with the sandbox.
+    ///
+    /// Potentially the runtime is executed in parallel and thus we need to wrap the registry in
+    /// `Arc<Mutex>` instead of `Rc<RefCell>`.
+    mock_registry: Arc<Mutex<MockRegistry<AccountId>>>,
+}
+
+impl<AccountId: Ord + Decode> InterceptingExtT for MockingExtension<AccountId> {
+    fn intercept_call(
+        &self,
+        contract_address: Vec<u8>,
+        _is_call: bool,
+        input_data: Vec<u8>,
+    ) -> Vec<u8> {
+        let contract_address = Decode::decode(&mut &contract_address[..])
+            .expect("Contract address should be decodable");
+
+        match self
+            .mock_registry
+            .lock()
+            .expect("Should be able to acquire registry")
+            .get(&contract_address)
+        {
+            // There is no mock registered for this address, so we return `None` to indicate that
+            // the call should be executed normally.
+            None => None::<()>.encode(),
+            // We intercept the call and return the result of the mock.
+            Some(mock) => {
+                let (selector, call_data) = input_data.split_at(4);
+                let selector: Selector = selector
+                    .try_into()
+                    .expect("Input data should contain at least selector bytes");
+
+                let result = mock
+                    .call(selector, call_data.to_vec())
+                    .expect("TODO: let the user define the fallback mechanism");
+
+                let result: ExecResult = Ok(ExecReturnValue {
+                    flags: ReturnFlags::empty(),
+                    data: result,
+                });
+
+                Some(result).encode()
+            }
+        }
+    }
 }

drink/src/mock.rs

@@ -0,0 +1,36 @@
+mod contract;
+mod error;
+mod mocking_api;
+
+use std::collections::BTreeMap;
+
+pub use contract::{mock_message, ContractMock, MessageMock, Selector};
+use error::MockingError;
+pub use mocking_api::MockingApi;
+
+/// Untyped result of a mocked call.
+pub type MockedCallResult = Result<Vec<u8>, MockingError>;
+
+/// A registry of mocked contracts.
+pub(crate) struct MockRegistry<AccountId: Ord> {
+    mocked_contracts: BTreeMap<AccountId, ContractMock>,
+}
+
+impl<AccountId: Ord> MockRegistry<AccountId> {
+    /// Creates a new registry.
+    pub fn new() -> Self {
+        Self {
+            mocked_contracts: BTreeMap::new(),
+        }
+    }
+
+    /// Registers `mock` for `address`. Returns the previous mock, if any.
+    pub fn register(&mut self, address: AccountId, mock: ContractMock) -> Option<ContractMock> {
+        self.mocked_contracts.insert(address, mock)
+    }
+
+    /// Returns the mock for `address`, if any.
+    pub fn get(&self, address: &AccountId) -> Option<&ContractMock> {
+        self.mocked_contracts.get(address)
+    }
+}

drink/src/mock/contract.rs

@@ -0,0 +1,64 @@
+use std::collections::BTreeMap;
+
+use parity_scale_codec::{Decode, Encode};
+
+use crate::{
+    errors::LangError,
+    mock::{error::MockingError, MockedCallResult},
+};
+
+/// Alias for a 4-byte selector.
+pub type Selector = [u8; 4];
+/// An untyped message mock.
+///
+/// Notice that in the end, we cannot operate on specific argument/return types. Rust won't let us
+/// have a collection of differently typed closures. Fortunately, we can assume that all types are
+/// en/decodable, so we can use `Vec<u8>` as a common denominator.
+pub type MessageMock = Box<dyn Fn(Vec<u8>) -> MockedCallResult + Send + Sync>;
+
+/// A contract mock.
+pub struct ContractMock {
+    messages: BTreeMap<Selector, MessageMock>,
+}
+
+impl ContractMock {
+    /// Creates a new mock without any message.
+    pub fn new() -> Self {
+        Self {
+            messages: BTreeMap::new(),
+        }
+    }
+
+    /// Adds a message mock.
+    pub fn with_message(mut self, selector: Selector, message: MessageMock) -> Self {
+        self.messages.insert(selector, message);
+        self
+    }
+
+    /// Try to call a message mock. Returns an error if there is no message mock for `selector`.
+    pub fn call(&self, selector: Selector, input: Vec<u8>) -> MockedCallResult {
+        match self.messages.get(&selector) {
+            None => Err(MockingError::MessageNotFound(selector)),
+            Some(message) => message(input),
+        }
+    }
+}
+
+impl Default for ContractMock {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// A helper function to create a message mock out of a typed closure.
+///
+/// In particular, it takes care of decoding the input and encoding the output. Also, wraps the
+/// return value in a `Result`, which is normally done implicitly by ink!.
+pub fn mock_message<Args: Decode, Ret: Encode, Body: Fn(Args) -> Ret + Send + Sync + 'static>(
+    body: Body,
+) -> MessageMock {
+    Box::new(move |encoded_input| {
+        let input = Decode::decode(&mut &*encoded_input).map_err(MockingError::ArgumentDecoding)?;
+        Ok(Ok::<Ret, LangError>(body(input)).encode())
+    })
+}