scrypto_test/environment/

env.rs

#![allow(clippy::test_attr_in_doctest)]

//! This module defines and implements the [`TestEnvironment`] struct.

use super::*;
use crate::prelude::*;

/// The environment that all tests of this testing framework are run against.
///
/// This struct may be thought of as the main struct in this testing framework which encapsulates a
/// a self-contained instance of the Radix Engine ([`EncapsulatedRadixEngine`]). The functionality
/// of the Radix Engine is exposed through the [`SystemApi`] which makes this testing environment no
/// less capable than Scrypto code.
///
/// ## Introduction
///
/// This testing framework is designed to allow you to write Scrypto-like code and use that to test
/// your packages and blueprints and follows a different approach from the `LedgerSimulator` class. The
/// test-runner is an in-memory ledger simulator which you can interact with as a user that submits
/// transactions to the network. The approach followed by this testing framework is different,
/// instead of submitting transactions, you're making invocations to the Radix Engine, getting
/// results back, and then writing assertions against what you got back.
///
/// Both the LedgerSimulator and this testing framework will prove to be useful throughout your blueprint
/// development journey. As an example, this testing framework allows you to disable some of kernel
/// modules that may get in your way when writing tests so it may be an optimal framework to use to
/// ensure that the "math checks out" in your blueprint code without needing to think about costing
/// or auth. However, when you're reaching the final stages of developing a blueprint you may want
/// tests that check that interactions with your blueprint will succeed in a simulated setting that
/// is close to the real setting, which is when the `LedgerSimulator` comes in. Overall, we may put these
/// two frameworks into two categories: This framework (named scrypto-test) is a framework for unit
/// testing your blueprints and is a good framework to use to check that your DeFi logic is correct.
/// The `LedgerSimulator` is an integration testing or an end-to-end testing framework to test that your
/// blueprints work in a simulated ledger with all of the costing limits, substate limits, and other
/// limits applied.
///
/// ## Features
///
/// This framework has many new features that developers may find useful when testing their packages
/// some of those features are:
///
/// * The ability to create mock [`Bucket`]s and [`Proof`]s through two main ways: by creating them
///   out of thin air, and by disabling the auth module and minting them. This functionality can be
///   found in the [`BucketFactory`] and [`ProofFactory`] structs and the [`CreationStrategy`].
/// * The ability to query the contents of [`Bucket`]s and [`Proof`]s for the purpose of writing
///   assertions against them. Not only that, but this testing framework allows you to call any
///   method you wish on these nodes. As an example, in a test, you can get a [`Bucket`] and then
///   create a proof out of it in manner similar to Scrypto.
/// * The ability to enable and disable kernel modules at runtime. The Radix Engine kernel is quite
///   modular with concepts such as auth, costing, and limits being implemented as kernel modules.
///   Disabling or enabling kernel modules at runtime can prove to be quite useful when writing DeFi
///   tests. As an example, you may want to not think about costing at all when writing tests and
///   thus you may opt to disable the costing module entirely and continue your test without it.
///   This can be done through [`TestEnvironment::disable_costing_module`].
/// * This testing framework uses test bindings to provide a higher-level API for calling methods
///   and functions on a blueprint without the need to do raw [`TestEnvironment::call_method_typed`]
///   or [`TestEnvironment::call_function_typed`]. The test bindings are generated by the blueprint
///   macro and are feature gated behind a `test` feature.
///
/// ## Getting Started
///
/// The following example shows a very simple test that gets XRD from the faucet and then asserts
/// that the amount is equal to what we expect.
///
/// ```
/// use scrypto_test::prelude::*;
///
/// fn my_test() -> Result<(), RuntimeError> {
///     // Arrange
///     let mut env = TestEnvironment::new();
///
///     // Act
///     let bucket = env.call_method_typed::<_, _, Bucket>(
///         FAUCET,
///         "free",
///         &(),
///     )?;
///
///     // Assert
///     let amount = bucket.amount(&mut env)?;
///     assert_eq!(amount, dec!("10000"));
///
///     Ok(())
/// }
/// ```
///
/// A few things to note about the code you see above:
///
/// * There is no transactions, worktop, receipt, manifests or anything of that sort! This part is
///   "not just hidden" from this testing framework but is actually non existent! The approach that
///   framework of wrapping a self-contained Radix Engine means that there is no need for manifests
///   or other transaction concepts.
/// * Methods such as [`Bucket::amount`] can be called to get the amount of resources in a bucket
///   and then assert against that.
///
/// ## Manipulating Kernel Modules
///
/// At runtime, the kernel modules can be enabled or disabled. For each kernel module there are four
/// methods on the [`TestEnvironment`]:
///
/// * A method to enable the kernel module (e.g., [`TestEnvironment::enable_costing_module`]).
/// * A method to disable the kernel module (e.g., [`TestEnvironment::disable_costing_module`]).
/// * A method to perform some action in a callback with the module enabled (e.g.,
///   [`TestEnvironment::with_costing_module_enabled`]).
/// * A method to perform some action in a callback with the module disabled (e.g.,
///   [`TestEnvironment::with_costing_module_disabled`]).
///
/// The simple enable and disable methods are quite straightforward: call them to enable or disable
/// a kernel module. The `with_*` methods are a little bit more intricate, they allow you to perform
/// some actions with a specific kernel either enabled or disabled and then resets the state of the
/// kernel modules afterwards. As an example:
///
/// ```
/// use scrypto_test::prelude::*;
///
/// // Arrange
/// let mut env = TestEnvironment::new();
///
/// // Act
/// let bucket = env.with_auth_module_disabled(|env| {
///     /* Auth Module is disabled just before this point */
///     ResourceManager(XRD).mint_fungible(100.into(), env)
///     /* Kernel modules are reset just after this point. */
/// })?;
///
/// // Assert
/// let amount = bucket.amount(&mut env)?;
/// assert_eq!(amount, dec!("100"));
///
/// # Ok::<(), RuntimeError>(())
/// ```
///
/// ## Common Arranges or Teardowns
///
/// There are cases where you may have many tests that all share a large portion of your arrange
/// or teardown logic. While this framework does not specifically provide solutions for this, there
/// are many useful Rust patterns that may be employed here to allow you to do this: the simplest
/// and the most elegant is probably by using callback functions.
///
/// Imagine this, you're building a Dex and many of the tests you write require you to have two
/// resources with a very large supply so you can write your tests with. You can achieve this by
/// doing something like:
///
/// ```
/// use scrypto_test::prelude::*;
///
/// pub fn two_resource_environment<F>(func: F)
/// where
///     F: FnOnce(DefaultTestEnvironment, FungibleBucket, FungibleBucket),
/// {
///     let mut env = TestEnvironment::new();
///     let bucket1 = ResourceBuilder::new_fungible(OwnerRole::None)
///         .mint_initial_supply(dec!("100000000000"), &mut env)
///         .unwrap();
///     let bucket2 = ResourceBuilder::new_fungible(OwnerRole::None)
///         .mint_initial_supply(dec!("100000000000"), &mut env)
///         .unwrap();
///     func(env, bucket1, bucket2)
///
///     /* Potential teardown happens here */
/// }
///
/// #[test]
/// fn contribution_provides_expected_amount_of_pool_units() {
///     two_resource_environment(|mut env, bucket1, bucket2| {
///         /* Your test goes here */
///     })
/// }
/// ```
///
/// You may have a function like `two_resource_environment` seen above which sets up the environment
/// and then some callback and potentially then executes some teardown code. Another way to do this
/// would be through simple factory and destructor methods.
/// ```
pub struct TestEnvironment<D>(pub(super) EncapsulatedRadixEngine<D>)
where
    D: SubstateDatabase + CommittableSubstateDatabase + 'static;

pub type DefaultTestEnvironment = TestEnvironment<InMemorySubstateDatabase>;

impl DefaultTestEnvironment {
    pub fn new() -> Self {
        TestEnvironmentBuilder::new().build()
    }
}

impl<D> TestEnvironment<D>
where
    D: SubstateDatabase + CommittableSubstateDatabase + 'static,
{
    // ===================
    // region:invocations
    // ===================

    /// Invokes a function on the provided blueprint and package with the given arguments.
    ///
    /// This method is a typed version of the [`SystemBlueprintApi::call_function`] which Scrypto
    /// encodes the arguments and Scrypto decodes the returns on behalf of the caller. This method
    /// assumes that the caller is correct about the argument and return types and panics if the
    /// encoding or decoding fails.
    ///
    /// # Arguments
    ///
    /// * `package_address`: [`PackageAddress`] - The address of the package that contains the
    ///   blueprint.
    /// * `blueprint_name`: [`&str`] - The name of the blueprint.
    /// * `function_name`: [`&str`] - The nae of the function.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    ///   is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    ///   was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    ///   [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    ///   [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    ///   the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    ///   could be decoded as the desired output type.
    pub fn call_function_typed<I, O>(
        &mut self,
        package_address: PackageAddress,
        blueprint_name: &str,
        function_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_function(package_address, blueprint_name, function_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    /// Invokes a method on the main module of a node with the provided typed arguments.
    ///
    /// This method is a typed version of the [`SystemObjectApi::call_method`] which Scrypto encodes
    /// the arguments and Scrypto decodes the returns on behalf of the caller. This method assumes
    /// that the caller is correct about the argument and return types and panics if the encoding or
    /// decoding fails.
    ///
    /// # Arguments
    ///
    /// * `node_id`: `T` - The node to invoke the method on. This is a generic argument that's
    ///   fulfilled by any type that implements [`Into<NodeId>`], thus, any address type can be used.
    /// * `method_name`: [`&str`] - The name of the method to invoke.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    ///   is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    ///   was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    ///   [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    ///   [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    ///   the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    ///   could be decoded as the desired output type.
    pub fn call_method_typed<N, I, O>(
        &mut self,
        node_id: N,
        method_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId>,
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_method(&node_id.into(), method_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    /// Invokes a method on the main module of a node with the provided typed arguments.
    ///
    /// This method is a typed version of the [`SystemObjectApi::call_method`] which Scrypto encodes
    /// the arguments and Scrypto decodes the returns on behalf of the caller. This method assumes
    /// that the caller is correct about the argument and return types and panics if the encoding or
    /// decoding fails.
    ///
    /// # Arguments
    ///
    /// * `node_id`: `T` - The node to invoke the method on. This is a generic argument that's
    ///   fulfilled by any type that implements [`Into<NodeId>`], thus, any address type can be used.
    /// * `method_name`: [`&str`] - The name of the method to invoke.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    ///   is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    ///   was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    ///   [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    ///   [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    ///   the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    ///   could be decoded as the desired output type.
    pub fn call_direct_access_method_typed<N, I, O>(
        &mut self,
        node_id: N,
        method_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId>,
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_direct_access_method(&node_id.into(), method_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    /// Invokes a method on a module of a node with the provided typed arguments.
    ///
    /// This method is a typed version of the [`SystemObjectApi::call_method`] which Scrypto encodes
    /// the arguments and Scrypto decodes the returns on behalf of the caller. This method assumes
    /// that the caller is correct about the argument and return types and panics if the encoding or
    /// decoding fails.
    ///
    /// # Arguments
    ///
    /// * `node_id`: `T` - The node to invoke the method on. This is a generic argument that's
    ///   fulfilled by any type that implements [`Into<NodeId>`], thus, any address type can be used.
    /// * `module`: [`AttachedModuleId`] - The module id.
    /// * `method_name`: [`&str`] - The name of the method to invoke.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    ///   is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    ///   was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    ///   [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    ///   [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    ///   the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    ///   could be decoded as the desired output type.
    pub fn call_module_method_typed<N, I, O>(
        &mut self,
        node_id: N,
        module: AttachedModuleId,
        method_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId>,
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_module_method(&node_id.into(), module, method_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    //=======================
    // endregion:invocations
    // region:kernel-access
    //=======================

    pub fn with_kernel<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&TestKernel<'_, D>) -> O,
    {
        self.0.with_kernel(callback)
    }

    pub fn with_kernel_mut<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut TestKernel<'_, D>) -> O,
    {
        self.0.with_kernel_mut(callback)
    }

    //=========================
    // endregion:kernel-access
    // region:kernel-modules
    //=========================

    /// Enables the kernel trace kernel module of the Radix Engine.
    pub fn enable_kernel_trace_module(&mut self) {
        self.enable_module(EnabledModules::KERNEL_TRACE)
    }

    /// Enables the limits kernel module of the Radix Engine.
    pub fn enable_limits_module(&mut self) {
        self.enable_module(EnabledModules::LIMITS)
    }

    /// Enables the costing kernel module of the Radix Engine.
    pub fn enable_costing_module(&mut self) {
        self.enable_module(EnabledModules::COSTING)
    }

    /// Enables the auth kernel module of the Radix Engine.
    pub fn enable_auth_module(&mut self) {
        self.enable_module(EnabledModules::AUTH)
    }

    /// Enables the transaction env kernel module of the Radix Engine.
    pub fn enable_transaction_runtime_module(&mut self) {
        self.enable_module(EnabledModules::TRANSACTION_RUNTIME)
    }

    /// Enables the execution trace kernel module of the Radix Engine.
    pub fn enable_execution_trace_module(&mut self) {
        self.enable_module(EnabledModules::EXECUTION_TRACE)
    }

    /// Disables the kernel trace kernel module of the Radix Engine.
    pub fn disable_kernel_trace_module(&mut self) {
        self.disable_module(EnabledModules::KERNEL_TRACE)
    }

    /// Disables the limits kernel module of the Radix Engine.
    pub fn disable_limits_module(&mut self) {
        self.disable_module(EnabledModules::LIMITS)
    }

    /// Disables the costing kernel module of the Radix Engine.
    pub fn disable_costing_module(&mut self) {
        self.disable_module(EnabledModules::COSTING)
    }

    /// Disables the auth kernel module of the Radix Engine.
    pub fn disable_auth_module(&mut self) {
        self.disable_module(EnabledModules::AUTH)
    }

    /// Disables the transaction env kernel module of the Radix Engine.
    pub fn disable_transaction_runtime_module(&mut self) {
        self.disable_module(EnabledModules::TRANSACTION_RUNTIME)
    }

    /// Disables the execution trace kernel module of the Radix Engine.
    pub fn disable_execution_trace_module(&mut self) {
        self.disable_module(EnabledModules::EXECUTION_TRACE)
    }

    /// Calls the passed `callback` with the kernel trace kernel module enabled and then resets the
    /// state of the kernel modules.
    pub fn with_kernel_trace_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_kernel_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the limits kernel module enabled and then resets the state
    /// of the kernel modules.
    pub fn with_limits_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_limits_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the costing kernel module enabled and then resets the state
    /// of the kernel modules.
    pub fn with_costing_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_costing_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the auth kernel module enabled and then resets the state of
    /// the kernel modules.
    pub fn with_auth_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_auth_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the transaction env kernel module enabled and then
    /// resets the state of the kernel modules.
    pub fn with_transaction_runtime_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_transaction_runtime_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the execution trace kernel module enabled and then resets
    /// the state of the kernel modules.
    pub fn with_execution_trace_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_execution_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the kernel trace kernel module disabled and then resets the
    /// state of the kernel modules.
    pub fn with_kernel_trace_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_kernel_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the limits kernel module disabled and then resets the state
    /// of the kernel modules.
    pub fn with_limits_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_limits_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the costing kernel module disabled and then resets the
    /// state of the kernel modules.
    pub fn with_costing_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_costing_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the auth kernel module disabled and then resets the state
    /// of the kernel modules.
    pub fn with_auth_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_auth_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the transaction env kernel module disabled and then
    /// resets the state of the kernel modules.
    pub fn with_transaction_runtime_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_transaction_runtime_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the execution trace kernel module disabled and then resets
    /// the state of the kernel modules.
    pub fn with_execution_trace_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_execution_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Returns the bit flags representing the currently enabled kernel modules.
    pub fn enabled_modules(&self) -> EnabledModules {
        self.0
            .with_kernel(|kernel| kernel.kernel_callback().modules.enabled_modules)
    }

    /// Sets the bit flags representing the enabled kernel modules.
    pub fn set_enabled_modules(&mut self, enabled_modules: EnabledModules) {
        self.0.with_kernel_mut(|kernel| {
            kernel.kernel_callback_mut().modules.enabled_modules = enabled_modules
        })
    }

    /// Enables specific kernel module(s).
    pub fn enable_module(&mut self, module: EnabledModules) {
        self.0.with_kernel_mut(|kernel| {
            kernel.kernel_callback_mut().modules.enabled_modules |= module
        })
    }

    /// Disables specific kernel module(s).
    pub fn disable_module(&mut self, module: EnabledModules) {
        self.0.with_kernel_mut(|kernel| {
            kernel.kernel_callback_mut().modules.enabled_modules &= !module
        })
    }

    //==========================
    // endregion:kernel-modules
    // region:state
    //==========================

    /// Reads the state of a component and allows for a callback to be executed over the decoded
    /// state.
    ///
    /// This method performs the steps needed to read the state of a component and then perform the
    /// various steps needed before the state can be read or used such as the locking, reading, and
    /// decoding of the substate and the various steps that need to be performed after the state is
    /// read such as unlocking the substate.
    ///
    /// Users of this method are expected to pass in a callback function to operate over the state
    /// as this is the main way that this method ensures that references do not escape out of this
    /// method after the substate is closed.
    ///
    /// # Arguments
    ///
    /// * `node_id`: [`N`] - The address of the component to read the state of. This is a generic
    ///   type parameter that's satisfied by any type that implements [`Into<NodeId>`].
    /// * `callback`: [`F`] - A callback function to call after the component state has been read
    ///   and decoded into the type specified by the generic parameter [`S`]. Anything returned from
    ///   this callback is returned from this method unless an error happens after the callback is
    ///   executed.
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The output of the callback function passed in or a runtime
    ///   error if one of the steps failed.
    ///
    /// # Panics
    ///
    /// This method panics if the component state can not be decoded as the generic type parameter
    /// [`S`].
    pub fn with_component_state<S, N, F, O>(
        &mut self,
        node_id: N,
        mut callback: F,
    ) -> Result<O, RuntimeError>
    where
        S: ScryptoDecode,
        N: Into<NodeId>,
        F: FnMut(&mut S, &mut Self) -> O,
    {
        let (handle, state) = self.0.with_kernel_mut(|kernel| {
            // Lock
            let handle = kernel.kernel_open_substate(
                &node_id.into(),
                MAIN_BASE_PARTITION,
                &SubstateKey::Field(ComponentField::State0.into()),
                LockFlags::read_only(),
                SystemLockData::Field(FieldLockData::Read),
            )?;

            // Read
            let state = kernel.kernel_read_substate(handle).map(|v| {
                let FieldSubstate::<ScryptoValue>::V1(FieldSubstateV1 { payload, .. }) =
                    v.as_typed().unwrap();
                scrypto_encode(&payload).unwrap()
            })?;

            Ok::<_, RuntimeError>((handle, state))
        })?;

        // Decode
        let mut state = scrypto_decode::<S>(&state).unwrap();

        // Callback
        let rtn = callback(&mut state, self);

        // Unlock
        self.0
            .with_kernel_mut(|kernel| kernel.kernel_close_substate(handle))?;

        Ok(rtn)
    }

    //========================
    // endregion:state
    // region:epoch&timestamp
    //========================

    /// Gets the current epoch from the Consensus Manager.
    pub fn get_current_epoch(&mut self) -> Epoch {
        Runtime::current_epoch(self).unwrap()
    }

    /// Sets the current epoch.
    pub fn set_current_epoch(&mut self, epoch: Epoch) {
        self.as_method_actor(
            CONSENSUS_MANAGER,
            ModuleId::Main,
            CONSENSUS_MANAGER_NEXT_ROUND_IDENT,
            |env| -> Result<(), RuntimeError> {
                let manager_handle = env
                    .actor_open_field(
                        ACTOR_STATE_SELF,
                        ConsensusManagerField::State.into(),
                        LockFlags::MUTABLE,
                    )
                    .unwrap();
                let mut manager_substate =
                    env.field_read_typed::<VersionedConsensusManagerState>(manager_handle)?;

                manager_substate.as_unique_version_mut().epoch = epoch;

                env.field_write_typed(manager_handle, &manager_substate)?;
                env.field_close(manager_handle)?;
                Ok(())
            },
        )
        .unwrap()
        .unwrap();
    }

    /// Gets the current time stamp from the Consensus Manager.
    pub fn get_current_time(&mut self) -> Instant {
        Runtime::current_time(TimePrecision::Second, self).unwrap()
    }

    pub fn set_current_time(&mut self, instant: Instant) {
        self.as_method_actor(
            CONSENSUS_MANAGER,
            ModuleId::Main,
            CONSENSUS_MANAGER_NEXT_ROUND_IDENT,
            |env| -> Result<(), RuntimeError> {
                let handle = env.actor_open_field(
                    ACTOR_STATE_SELF,
                    ConsensusManagerField::ProposerMinuteTimestamp.into(),
                    LockFlags::MUTABLE,
                )?;
                let mut proposer_minute_timestamp =
                    env.field_read_typed::<ConsensusManagerProposerMinuteTimestampFieldPayload>(
                        handle,
                    )?;
                proposer_minute_timestamp
                    .as_unique_version_mut()
                    .epoch_minute =
                    i32::try_from(instant.seconds_since_unix_epoch / 60).map_err(|_| {
                        RuntimeError::ApplicationError(ApplicationError::ConsensusManagerError(
                            ConsensusManagerError::InvalidProposerTimestampUpdate {
                                from_millis: env.get_current_time().seconds_since_unix_epoch * 1000,
                                to_millis: instant.seconds_since_unix_epoch * 1000,
                            },
                        ))
                    })?;
                env.field_write_typed(handle, &proposer_minute_timestamp)?;
                env.field_close(handle)?;

                let handle = env.actor_open_field(
                    ACTOR_STATE_SELF,
                    ConsensusManagerField::ProposerMilliTimestamp.into(),
                    LockFlags::MUTABLE,
                )?;
                let mut proposer_milli_timestamp =
                    env.field_read_typed::<ConsensusManagerProposerMilliTimestampFieldPayload>(
                        handle,
                    )?;
                proposer_milli_timestamp.as_unique_version_mut().epoch_milli =
                    instant.seconds_since_unix_epoch * 1000;
                env.field_write_typed(handle, &proposer_milli_timestamp)?;
                env.field_close(handle)?;

                Ok(())
            },
        )
        .unwrap()
        .unwrap();
    }

    //===========================
    // endregion:epoch&timestamp
    // region:helpers
    //===========================

    /// Allows us to perform some action as another actor.
    ///
    /// This function pushes a new call-frame onto the stack with the actor information we desire,
    /// performs the call-back, and then pops the call-frame off.
    pub(crate) fn as_method_actor<N, F, O>(
        &mut self,
        node_id: N,
        module_id: ModuleId,
        method_ident: &str,
        callback: F,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId> + Copy,
        F: FnOnce(&mut Self) -> O,
        O: ScryptoEncode,
    {
        let object_info = self.0.with_kernel_mut(|kernel| {
            SystemService::new(kernel).get_object_info(&node_id.into())
        })?;
        let auth_zone = self.0.with_kernel_mut(|kernel| {
            let mut system_service = SystemService::new(kernel);
            AuthModule::on_call_fn_mock(
                &mut system_service,
                Some((&node_id.into(), false)),
                Default::default(),
                Default::default(),
            )
        })?;
        let actor = Actor::Method(MethodActor {
            method_type: match module_id {
                ModuleId::Main => MethodType::Main,
                ModuleId::Royalty => MethodType::Module(AttachedModuleId::Royalty),
                ModuleId::Metadata => MethodType::Module(AttachedModuleId::Metadata),
                ModuleId::RoleAssignment => MethodType::Module(AttachedModuleId::RoleAssignment),
            },
            ident: method_ident.to_owned(),
            node_id: node_id.into(),
            auth_zone,
            object_info,
        });
        self.as_actor(actor, callback)
    }

    /// Allows us to perform some action as another actor.
    ///
    /// This function pushes a new call-frame onto the stack with the actor information we desire,
    /// performs the call-back, and then pops the call-frame off.
    pub(crate) fn as_actor<F, O>(&mut self, actor: Actor, callback: F) -> Result<O, RuntimeError>
    where
        F: FnOnce(&mut Self) -> O,
        O: ScryptoEncode,
    {
        // Creating the next frame.
        let mut message =
            CallFrameMessage::from_input(&IndexedScryptoValue::from_typed(&()), &actor);
        self.0.with_kernel_mut(|kernel| {
            let (substate_io, current_frame) = kernel.kernel_current_frame_mut();
            message
                .copy_global_references
                .extend(current_frame.stable_references().keys());
            let new_frame =
                CallFrame::new_child_from_parent(substate_io, current_frame, actor, message)
                    .unwrap();
            let old = core::mem::replace(current_frame, new_frame);
            kernel.kernel_prev_frame_stack_mut().push(old);
        });

        // Executing the callback
        let rtn = callback(self);

        // Constructing the message from the returns
        let message = {
            let indexed_scrypto_value = IndexedScryptoValue::from_typed(&rtn);
            CallFrameMessage {
                move_nodes: indexed_scrypto_value.owned_nodes().clone(),
                copy_global_references: indexed_scrypto_value.references().clone(),
                ..Default::default()
            }
        };

        // Popping the last frame & Passing message
        self.0
            .with_kernel_mut(|kernel| -> Result<(), RuntimeError> {
                let mut previous_frame = kernel.kernel_prev_frame_stack_mut().pop().unwrap();
                let (substate_io, current_frame) = kernel.kernel_current_frame_mut();

                CallFrame::pass_message(
                    substate_io,
                    current_frame,
                    &mut previous_frame,
                    message.clone(),
                )
                .map_err(CallFrameError::PassMessageError)
                .map_err(KernelError::CallFrameError)?;

                *current_frame = previous_frame;
                Ok(())
            })?;

        Ok(rtn)
    }

    //===================
    // endregion:helpers
    //===================
}

impl Default for TestEnvironment<InMemorySubstateDatabase> {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    pub fn test_env_can_be_created() {
        let _ = TestEnvironment::new();
    }

    #[test]
    pub fn test_time_set() {
        // Arrange
        let mut env = TestEnvironment::new();
        let mut current_time = env.get_current_time();
        current_time.seconds_since_unix_epoch += 60; // add 1 minute

        // Act
        env.set_current_time(current_time);

        // Assert
        let t1 = Runtime::current_time(TimePrecision::Second, &mut env).unwrap();
        let t2 = Runtime::current_time(TimePrecision::Minute, &mut env).unwrap();

        assert_eq!(t1, t2)
    }
}