o2-tools 0.1.10

use crate::{
    CallOption,
    call_data,
    call_handler_ext::CallHandlerExt,
    order_book::{
        CreateOrderParams,
        OrderBookManager,
    },
    signature_ext::{
        DomainExt,
        SRC16Encode,
    },
    trade_account_deploy::{
        CallContractArg,
        CallParams,
        Domain,
        SRC16Domain,
        Secp256k1,
        Session,
        SessionArgs,
        Signature,
        State,
        Time,
        TradeAccountDeploy,
        TradingAccount,
        TradingAccountProxy,
    },
};
use fuel_core_types::fuel_types::ChainId;
use fuels::{
    accounts::{
        signers::private_key::PrivateKeySigner,
        wallet::Unlocked,
    },
    core::{
        codec::calldata,
        traits::Tokenizable,
    },
    crypto::Message,
    prelude::*,
    types::{
        Bytes32,
        Identity,
    },
};
use futures::FutureExt;
use rand::{
    SeedableRng,
    rngs::StdRng,
};
use std::time::{
    SystemTime,
    UNIX_EPOCH,
};

pub const TRADE_ACCOUNT_DOMAIN: &str = "TradeAccount";
pub const TRADE_ACCOUNT_PROXY_DOMAIN: &str = "TradeAccountProxy";
pub const TRADE_ACCOUNT_VERSION: &str = "1";
pub const TRADE_ACCOUNT_PROXY_VERSION: &str = "1";

/// Parameters for creating a contract call that will be signed by a trading session.
/// Used to prepare the data that needs to be hashed and signed for session-based calls.
pub struct CallContractParams {
    /// The contract ID to call
    pub contract_id: ContractId,
    /// The function selector (method signature) to call
    pub function_selector: Vec<u8>,
    /// Optional parameters for forwarding assets with the call
    pub forward: CallParams,
    /// Optional encoded arguments for the function call
    pub args: Option<Vec<u8>>,
    /// Current nonce to prevent replay attacks
    pub nonce: u64,
    /// Variable outputs
    pub variable_outputs: Option<u16>,
}

/// Represents a contract call that can be made within a trading session.
#[derive(Debug, Clone)]
pub enum ContractCalls {
    Call(CallContractArgs),
    Calls(CallContractsArgs),
}

impl ContractCalls {
    pub fn contract_ids(&self) -> Vec<ContractId> {
        match self {
            ContractCalls::Call(call) => vec![call.call.contract_id],
            ContractCalls::Calls(calls) => {
                calls.calls.iter().map(|c| c.contract_id).collect()
            }
        }
    }
}

/// Signed arguments ready to be submitted for a session-based contract call.
/// Contains the signature and all parameters needed to execute the call.
#[derive(Debug, Clone)]
pub struct CallContractArgs {
    /// Digital signature authorizing the call
    pub signature: Signature,
    /// The contract ID to call
    pub call: CallContractArg,
    /// Variable outputs
    pub variable_outputs: Option<u16>,
    /// Other contracts being called
    pub contracts: Option<Vec<ContractId>>,
}

/// Signed arguments ready to be submitted for a session-based contract call.
/// Contains the signature and all parameters needed to execute the call.
#[derive(Debug, Clone)]
pub struct CallContractsArgs {
    /// Digital signature authorizing the call
    pub signature: Signature,
    /// The contract ID to call
    pub calls: Vec<CallContractArg>,
    /// Variable outputs
    pub variable_outputs: Option<u16>,
    /// Other contracts being called
    pub contracts: Option<Vec<ContractId>>,
}

impl CallContractParams {
    /// Creates new call contract parameters.
    ///
    /// # Arguments
    /// * `contract_id` - The contract ID to call
    /// * `function_selector` - The function selector (method signature) bytes
    /// * `forward` - Optional call parameters for forwarding assets
    /// * `args` - Optional encoded arguments for the function call
    /// * `nonce` - Current nonce to prevent replay attacks
    pub fn new(
        contract_id: ContractId,
        function_selector: Vec<u8>,
        forward: CallParams,
        args: Option<Vec<u8>>,
        nonce: u64,
        variable_outputs: Option<u16>,
    ) -> Self {
        Self {
            contract_id,
            function_selector,
            forward,
            args,
            nonce,
            variable_outputs,
        }
    }

    pub fn call_args(&self) -> CallContractArg {
        let function_selector = Bytes(self.function_selector.clone());
        let call_params = CallParams {
            coins: self.forward.coins,
            asset_id: self.forward.asset_id,
            gas: self.forward.gas,
        };

        CallContractArg {
            contract_id: self.contract_id,
            function_selector: function_selector.clone(),
            call_params: call_params.clone(),
            call_data: self.args.clone().map(Bytes),
        }
    }

    /// Computes SHA256 hash of all call parameters for signing.
    /// This hash is used to create a digital signature that authorizes the call.
    ///
    /// # Returns
    /// * `Some(Message)` - The SHA256 hash of the parameters
    /// * `None` - If hashing fails (currently always returns Some)
    pub fn message(&self) -> Message {
        let call_contract = self.call_args();
        generate_session_signing_payload(self.nonce, call_contract)
    }
}

pub fn generate_session_signing_payload<T: Tokenizable>(
    nonce: u64,
    call_contract_arg: T,
) -> Message {
    Message::new(call_data!(nonce, call_contract_arg))
}

/// High-level interface for managing a trade account with session-based authorization.
/// Provides methods for creating sessions, making signed calls, and managing the account.
#[derive(Clone)]
pub struct TradeAccountManager<Wallet: Account + Clone> {
    /// The wallet that owns this trade account
    pub owner: Wallet,
    /// The proxy contract instance for this trade account
    pub proxy: TradingAccountProxy<Wallet>,
    /// The trade account contract instance
    pub contract: TradingAccount<Wallet>,
    /// Private key signer for creating trading sessions
    pub session_signer: Option<PrivateKeySigner>,
    /// The nonce for the trade account
    pub nonce: u64,
}

impl TradeAccountManager<Wallet> {
    pub fn contract_id(&self) -> ContractId {
        self.contract.id()
    }

    pub fn identity(&self) -> Identity {
        Identity::ContractId(self.contract.contract_id())
    }

    pub fn session_signer(&self) -> anyhow::Result<PrivateKeySigner> {
        self.session_signer
            .clone()
            .ok_or_else(|| anyhow::anyhow!("Session signer not initialized"))
    }

    pub async fn new(owner: &Wallet, contract_id: ContractId) -> anyhow::Result<Self> {
        let mut trade_account = Self::new_with_nonce(owner, contract_id, 0);
        trade_account.fetch_nonce().await?;
        Ok(trade_account)
    }

    pub fn new_with_nonce(owner: &Wallet, contract_id: ContractId, nonce: u64) -> Self {
        let proxy = TradingAccountProxy::new(contract_id, owner.clone());
        let contract = TradingAccount::new(contract_id, owner.clone());
        Self {
            owner: owner.clone(),
            contract,
            proxy,
            session_signer: None,
            nonce,
        }
    }

    /// Creates a new TradeAccount instance from a deployed trade account.
    /// Sets up a session signer and establishes a trading session.
    ///
    /// # Arguments
    /// * `owner` - The wallet that owns the trade account
    /// * `trade_account_deploy` - The deployment information containing proxy and contract details
    ///
    /// # Returns
    /// * `Ok(TradeAccount)` - A configured trade account ready for use
    /// * `Err(anyhow::Error)` - If creation or session setup fails
    pub fn create(
        owner: &Wallet,
        trade_account_deploy: &TradeAccountDeploy<Wallet>,
    ) -> anyhow::Result<Self, anyhow::Error> {
        if trade_account_deploy.proxy_id.is_none() || trade_account_deploy.proxy.is_none()
        {
            return Err(anyhow::anyhow!(
                "Trade account deploy proxy ID or instance is not set"
            ));
        }
        let trade_account: TradingAccount<Wallet> =
            TradingAccount::new(trade_account_deploy.proxy_id.unwrap(), owner.clone());
        let trade_account_instance = Self {
            owner: owner.clone(),
            contract: trade_account,
            session_signer: None,
            proxy: trade_account_deploy.proxy.clone().unwrap(),
            nonce: 0,
        };
        Ok(trade_account_instance)
    }

    /// Creates a new TradeAccount instance from a deployed trade account.
    /// Sets up a session signer and establishes a trading session.
    ///
    /// # Arguments
    /// * `owner` - The wallet that owns the trade account
    /// * `trade_account_deploy` - The deployment information containing proxy and contract details
    ///
    /// # Returns
    /// * `Ok(TradeAccount)` - A configured trade account ready for use
    /// * `Err(anyhow::Error)` - If creation or session setup fails
    pub async fn create_with_session(
        fee_payer: &Wallet,
        owner: &Wallet,
        contract_ids: &[ContractId],
        trade_account_deploy: &TradeAccountDeploy<Wallet>,
        call_option: CallOption,
    ) -> anyhow::Result<Self, anyhow::Error> {
        let mut trade_account_instance = Self::create(owner, trade_account_deploy)?;
        let mut rng = StdRng::seed_from_u64(2322u64);
        let session_signer = PrivateKeySigner::random(&mut rng);
        let session_address = Identity::Address(session_signer.address());
        trade_account_instance.session_signer = Some(session_signer.clone());
        let session =
            trade_account_instance.new_session(session_address, contract_ids, None);

        let chain_id = owner.provider().consensus_parameters().await?.chain_id();

        let signature = create_personal_signature_args(
            chain_id,
            owner,
            trade_account_instance.nonce,
            Some(Some(session.clone())),
            String::from("set_session"),
        );

        let mut call_handler = trade_account_instance
            .contract
            .methods()
            .set_session(Some(signature), Some(session));
        call_handler.account = fee_payer.clone();

        match call_option {
            CallOption::AwaitBlock => {
                call_handler.call().await?;
                trade_account_instance.fetch_nonce().await?;
            }
            CallOption::AwaitPreconfirmation(ops) => {
                call_handler
                    .almost_sync_call(
                        &ops.data_builder,
                        &ops.utxo_manager,
                        &ops.tx_config,
                    )
                    .await?;
                trade_account_instance.increment_nonce();
            }
        }

        Ok(trade_account_instance)
    }

    /// Creates a new trading session with the specified address.
    /// Sessions allow temporary delegation of trading permissions with a time-based expiry.
    ///
    /// # Arguments
    /// * `session_address` - The identity that will be authorized to trade
    ///
    /// # Returns
    /// * `Session` - A new session with 30 days expiry
    pub fn new_session(
        &self,
        session_address: Identity,
        contract_ids: &[ContractId],
        expiry: Option<u64>,
    ) -> Session {
        let default_expiry = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("Time went backwards")
            .as_millis() as u64
            + (30 * 24 * 60 * 60 * 1000); // 30 days in milliseconds
        Session {
            session_id: session_address,
            expiry: Time {
                unix: expiry.unwrap_or(default_expiry),
            },
            contract_ids: contract_ids.to_vec(),
        }
    }

    /// Creates a new trading session for typed session arguments
    pub fn new_session_args(
        &self,
        nonce: u64,
        session_address: Identity,
        contract_ids: &[ContractId],
        expiry: Option<u64>,
    ) -> SessionArgs {
        let default_expiry = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("Time went backwards")
            .as_millis() as u64
            + (30 * 24 * 60 * 60 * 1000); // 30 days in milliseconds
        SessionArgs {
            nonce,
            session_id: session_address,
            expiry: Time {
                unix: expiry.unwrap_or(default_expiry),
            },
            contract_ids: contract_ids.to_vec(),
        }
    }

    /// Gets the current nonce for this trade account.
    /// Nonces are used to prevent replay attacks in session-based calls.
    ///
    /// # Returns
    /// * `Ok(u64)` - The current nonce value
    /// * `Err(anyhow::Error)` - If the contract call fails
    pub async fn fetch_nonce(&mut self) -> anyhow::Result<u64> {
        let nonce = self
            .contract
            .methods()
            .get_nonce()
            .simulate(Execution::state_read_only())
            .await?
            .value;
        self.nonce = nonce;
        Ok(nonce)
    }

    pub fn increment_nonce(&mut self) -> u64 {
        self.nonce += 1;
        self.nonce
    }

    /// Gets the address of the trade account owner from the proxy contract.
    ///
    /// # Returns
    /// * `Address` - The owner's address
    ///
    /// # Panics
    /// * If the proxy owner is not initialized
    pub async fn owner_address(&self) -> Address {
        let state: State = self
            .proxy
            .methods()
            .proxy_owner()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;
        if let State::Initialized(address) = state {
            match address {
                Identity::Address(address) => address,
                Identity::ContractId(contract_id) => Address::new(contract_id.into()),
            }
        } else {
            panic!("Proxy owner is not initialized");
        }
    }

    /// Creates signed call arguments for a session-based contract call.
    /// Takes call parameters, hashes them, and creates a digital signature using the session signer.
    ///
    /// # Arguments
    /// * `params` - The call parameters to sign (contract ID, function, args, nonce, etc.)
    ///
    /// # Returns
    /// * `Ok(CallContractArgs)` - Signed arguments ready for submission
    /// * `Err(anyhow::Error)` - If hashing or signing fails
    pub async fn create_call_args(
        &self,
        params: CallContractParams,
    ) -> anyhow::Result<CallContractArgs, anyhow::Error> {
        let signature = self.session_signer()?.sign(params.message()).await?;
        let signature = Signature::Secp256k1(Secp256k1 {
            bits: signature.as_slice().try_into()?,
        });
        Ok(CallContractArgs {
            signature,
            variable_outputs: params.variable_outputs,
            call: params.call_args(),
            contracts: None,
        })
    }

    /// Creates signed call arguments for a session-based contract call.
    /// Takes call parameters, hashes them, and creates a digital signature using the session signer.
    ///
    /// # Arguments
    /// * `params` - The call parameters to sign (contract ID, function, args, nonce, etc.)
    ///
    /// # Returns
    /// * `Ok(CallContractArgs)` - Signed arguments ready for submission
    /// * `Err(anyhow::Error)` - If hashing or signing fails
    pub async fn create_calls_args(
        &self,
        params: Vec<CallContractParams>,
    ) -> anyhow::Result<CallContractsArgs, anyhow::Error> {
        let call_contract_params = params
            .iter()
            .map(|param| param.call_args())
            .collect::<Vec<_>>();
        let message =
            generate_session_signing_payload(self.nonce, call_contract_params.clone());
        let signature = self.session_signer()?.sign(message).await?;
        let signature = Signature::Secp256k1(Secp256k1 {
            bits: signature.as_slice().try_into()?,
        });
        Ok(CallContractsArgs {
            signature,
            calls: call_contract_params,
            variable_outputs: None,
            contracts: None,
        })
    }

    pub fn call_contract(
        &self,
        call_args: &CallContractArgs,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = call_args.contracts.clone().unwrap_or_default();
        self.contract
            .methods()
            .call_contract(Some(call_args.signature.clone()), call_args.call.clone())
            .with_contract_ids(&contracts)
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub fn call_contracts(
        &self,
        call_args: &CallContractsArgs,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = call_args.contracts.clone().unwrap_or_default();
        self.contract
            .methods()
            .call_contracts(Some(call_args.signature.clone()), call_args.calls.clone())
            .with_contract_ids(&contracts)
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub fn session_call_contract(
        &self,
        call_args: &CallContractArgs,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        self.contract
            .methods()
            .session_call_contract(call_args.signature.clone(), call_args.call.clone())
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub fn session_call_contracts(
        &self,
        call_args: &CallContractsArgs,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        self.contract
            .methods()
            .session_call_contracts(call_args.signature.clone(), call_args.calls.clone())
    }

    pub fn cancel_orders_args(
        &self,
        order_book: &OrderBookManager<Wallet>,
        order_id: Bytes32,
        gas: Option<u64>,
    ) -> CallContractParams {
        CallContractParams::new(
            order_book.contract.contract_id(),
            crate::fn_selector!(cancel_order(OrderId)),
            CallParams {
                coins: 0,
                asset_id: AssetId::default(),
                gas: gas.unwrap_or(u64::MAX),
            },
            Some(call_data!(*order_id)),
            self.nonce,
            None,
        )
    }

    pub async fn cancel_order(
        &mut self,
        order_book: &OrderBookManager<Wallet>,
        order_id: Bytes32,
        gas: Option<u64>,
    ) -> anyhow::Result<CallContractArgs, anyhow::Error> {
        let call_args = self
            .create_call_args(self.cancel_orders_args(order_book, order_id, gas))
            .await?;
        self.increment_nonce();
        Ok(call_args)
    }

    pub async fn cancel_orders(
        &mut self,
        order_book: &OrderBookManager<Wallet>,
        order_ids: &[Bytes32],
        gas: Option<u64>,
    ) -> anyhow::Result<CallContractsArgs, anyhow::Error> {
        let call_contract_params = order_ids
            .iter()
            .map(|order_id| self.cancel_orders_args(order_book, *order_id, gas))
            .collect::<Vec<_>>();
        let call_contracts_args = self.create_calls_args(call_contract_params).await?;
        self.increment_nonce();
        Ok(call_contracts_args)
    }

    pub fn create_order_args(
        &self,
        order_book: &OrderBookManager<Wallet>,
        params: &CreateOrderParams,
        gas: Option<u64>,
    ) -> CallContractParams {
        CallContractParams::new(
            order_book.contract.contract_id(),
            crate::fn_selector!(create_order(OrderArgs)),
            order_book.create_call_params(params, gas),
            Some(call_data!(params.to_order_args())),
            self.nonce,
            None,
        )
    }

    pub async fn create_order(
        &mut self,
        order_book: &OrderBookManager<Wallet>,
        params: &CreateOrderParams,
        gas: Option<u64>,
    ) -> anyhow::Result<CallContractArgs, anyhow::Error> {
        let call_args = self
            .create_call_args(self.create_order_args(order_book, params, gas))
            .await?;
        self.increment_nonce();
        Ok(call_args)
    }

    pub async fn create_orders(
        &mut self,
        order_book: &OrderBookManager<Wallet>,
        params: &[CreateOrderParams],
        gas: Option<u64>,
    ) -> anyhow::Result<CallContractsArgs, anyhow::Error> {
        let call_contract_params = params
            .iter()
            .map(|param| self.create_order_args(order_book, param, gas))
            .collect::<Vec<_>>();
        let call_contracts_args = self.create_calls_args(call_contract_params).await?;
        self.increment_nonce();
        Ok(call_contracts_args)
    }
}

impl From<SessionArgs> for Session {
    fn from(args: SessionArgs) -> Self {
        Self {
            session_id: args.session_id,
            expiry: args.expiry,
            contract_ids: args.contract_ids,
        }
    }
}

pub fn generate_signing_payload<T>(
    nonce: u64,
    chain_id: u64,
    f_name: String,
    args: Option<T>,
) -> Message
where
    T: Tokenizable,
{
    // Fuel message prefix.
    // This is used to create a message that can be signed by the Fuel address.
    // The prefix is `b"Fuel Signed Message:\n"`.
    let mut message_bytes: Vec<u8> = vec![
        25, 70, 117, 101, 108, 32, 83, 105, 103, 110, 101, 100, 32, 77, 101, 115, 115,
        97, 103, 101, 58, 10,
    ];
    let mut bytes = if let Some(args_some) = args {
        calldata!((nonce, chain_id, f_name, args_some)).unwrap()
    } else {
        calldata!(nonce, chain_id, f_name).unwrap()
    };

    let mut num_bytes = bytes.len().to_string().into_bytes();

    message_bytes.append(&mut num_bytes);
    message_bytes.append(&mut bytes);

    Message::new(&message_bytes)
}

pub fn create_personal_signature_args<T>(
    chain_id: ChainId,
    wallet: &Wallet<Unlocked<PrivateKeySigner>>,
    nonce: u64,
    args: Option<T>,
    f_name: String,
) -> Signature
where
    T: Tokenizable,
{
    let message = generate_signing_payload(nonce, *chain_id, f_name, args);

    let fuel_signature = wallet
        .signer()
        .sign(message)
        .now_or_never()
        .expect("We use private key signer, so signing is immediate")
        .unwrap();
    Signature::Secp256k1(Secp256k1 {
        bits: *fuel_signature,
    })
}

pub async fn create_typed_signature<T>(
    wallet: &Wallet<Unlocked<PrivateKeySigner>>,
    args: T,
    trade_account_version: String,
    chain_id: u64,
    contract_id: ContractId,
) -> Signature
where
    T: SRC16Encode,
{
    let domain = Domain::SRC16Domain(SRC16Domain {
        name: Some(TRADE_ACCOUNT_DOMAIN.to_string()),
        version: Some(trade_account_version),
        chain_id: Some(chain_id.into()),
        verifying_contract: Some(contract_id),
        salt: None,
    });

    let message = Message::from_bytes(domain.encode(args));
    let fuel_signature = wallet.signer().sign(message).await.unwrap();
    Signature::Secp256k1(Secp256k1 {
        bits: *fuel_signature,
    })
}

pub async fn create_proxy_typed_signature<T>(
    wallet: &Wallet<Unlocked<PrivateKeySigner>>,
    args: T,
    trade_account_proxy_version: String,
    chain_id: u64,
    contract_id: ContractId,
) -> Signature
where
    T: SRC16Encode,
{
    let domain = Domain::SRC16Domain(SRC16Domain {
        name: Some(TRADE_ACCOUNT_PROXY_DOMAIN.to_string()),
        version: Some(trade_account_proxy_version),
        chain_id: Some(chain_id.into()),
        verifying_contract: Some(contract_id),
        salt: None,
    });
    let message = Message::from_bytes(domain.encode(args));
    let fuel_signature = wallet.signer().sign(message).await.unwrap();
    Signature::Secp256k1(Secp256k1 {
        bits: *fuel_signature,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        test_contracts::setup_large_contract,
        trade_account_deploy::{
            DeployConfig,
            TradeAccountDeployConfig,
        },
    };
    use fuels::test_helpers::{
        WalletsConfig,
        launch_custom_provider_and_get_wallets,
    };

    #[tokio::test]
    async fn test_trade_account() {
        // Start fuel-core
        let mut wallets = launch_custom_provider_and_get_wallets(
            WalletsConfig::new(Some(4), Some(1), Some(1_000_000_000_000)),
            None,
            None,
        )
        .await
        .unwrap();
        let deployer_wallet = wallets.pop().unwrap();
        let user_wallet = wallets.pop().unwrap();
        let receiver_wallet = wallets.pop().unwrap();
        let gaspayer_wallet = wallets.pop().unwrap();
        let (large_contract, large_contract_id) =
            setup_large_contract(&user_wallet).await;
        let deploy_config = DeployConfig::Latest(TradeAccountDeployConfig::default());
        let deployment = TradeAccountDeploy::deploy(&deployer_wallet, &deploy_config)
            .await
            .unwrap()
            .deploy_with_account(
                &user_wallet.address().into(),
                &deploy_config,
                &CallOption::AwaitBlock,
            )
            .await
            .unwrap();
        let trade_account = TradeAccountManager::create_with_session(
            &gaspayer_wallet,
            &user_wallet,
            &[large_contract_id],
            &deployment,
            CallOption::AwaitBlock,
        )
        .await
        .unwrap();
        let session_id =
            Identity::Address(trade_account.session_signer().unwrap().address());
        let balance_amount = 100_000_000u64;

        let is_valid = trade_account
            .contract
            .methods()
            .validate_session(session_id)
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;
        assert!(is_valid);

        let _ = user_wallet
            .force_transfer_to_contract(
                trade_account.contract.contract_id(),
                balance_amount,
                AssetId::default(),
                TxPolicies::default(),
            )
            .await
            .unwrap();
        let balances = trade_account.contract.get_balances().await.unwrap();
        let balance = balances.get(&AssetId::default()).unwrap();

        assert_eq!(user_wallet.address(), trade_account.owner_address().await);
        assert_eq!(*balance, balance_amount);

        let _ = trade_account
            .contract
            .methods()
            .withdraw(
                None,
                Identity::Address(receiver_wallet.address()),
                10_000u64,
                AssetId::default(),
            )
            .with_variable_output_policy(VariableOutputPolicy::Exactly(1))
            .call()
            .await
            .unwrap();

        assert_eq!(
            receiver_wallet
                .get_asset_balance(&AssetId::default())
                .await
                .unwrap(),
            1_000_000_010_000u128
        );

        let call_params = CallContractParams::new(
            large_contract_id,
            crate::fn_selector!(push_storage(u16)),
            CallParams::new(0, AssetId::default(), 100_000),
            Some(crate::call_data!(1u16)),
            trade_account.nonce,
            None,
        );
        let call_args = trade_account.create_call_args(call_params).await.unwrap();
        let result = trade_account
            .contract
            .clone()
            .with_account(gaspayer_wallet)
            .methods()
            .session_call_contract(call_args.signature, call_args.call)
            .with_contracts(&[&large_contract, &trade_account.contract])
            .with_contract_ids(&[large_contract_id, trade_account.contract.id()])
            .call()
            .await;

        assert!(result.is_ok());
    }
}