o2-tools 0.3.21-rc

use crate::{
    CallOption,
    call_handler_ext::CallHandlerExt,
    trade_account_deploy::{
        DeployConfig,
        State as TradingAccountState,
        TRADE_ACCOUNT_PROXY_BYTECODE,
        TradingAccountProxyConfigurables,
    },
    trial_trade_account_deploy::{
        TRIAL_TRADE_ACCOUNT_PROXY_BYTECODE,
        TRIAL_TRADE_ACCOUNT_PROXY_STORAGE,
        TrialTradingAccountProxyConfigurables,
    },
};
use anyhow::Result;
use fuels::{
    prelude::*,
    programs::calls::CallParameters,
    tx::StorageSlot,
    types::{
        Bits256,
        Bytes32,
        ContractId,
        Identity,
    },
};

abigen!(
    Contract(
        name = "TradeAccountRegistry",
        abi = "artifacts/trade-account-registry/trade-account-registry-abi.json"
    ),
    Contract(
        name = "TradeAccountRegistryProxy",
        abi = "artifacts/trade-account-registry-proxy/trade-account-registry-proxy-abi.json"
    )
);

pub const TRADE_ACCOUNT_REGISTER_BYTECODE: &[u8] =
    include_bytes!("../artifacts/trade-account-registry/trade-account-registry.bin");
pub const TRADE_ACCOUNT_REGISTER_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trade-account-registry/trade-account-registry-storage_slots.json"
);
pub const TRADE_ACCOUNT_REGISTER_PROXY_BYTECODE: &[u8] = include_bytes!(
    "../artifacts/trade-account-registry-proxy/trade-account-registry-proxy.bin"
);
pub const TRADE_ACCOUNT_REGISTER_PROXY_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trade-account-registry-proxy/trade-account-registry-proxy-storage_slots.json"
);

/// Configuration for deploying TradeAccountRegistry contracts.
/// Contains bytecode and storage slot information needed for deployment.
#[derive(Clone)]
pub struct TradeAccountRegistryDeployConfig {
    /// Bytecode for the TradeAccountRegistry contract
    pub registry_bytecode: Vec<u8>,
    /// Storage slots configuration for the TradeAccountRegistry contract
    pub registry_storage_slots: Vec<StorageSlot>,
    /// Configurables
    pub registry_config: TradeAccountRegistryConfigurables,
    /// Bytecode for the TradeAccountRegistryProxy contract
    pub registry_proxy_bytecode: Vec<u8>,
    /// Storage slots configuration for the TradeAccountRegistryProxy
    pub registry_proxy_storage_slots: Vec<StorageSlot>,
    /// Register proxy configurations
    pub registry_proxy_config: TradeAccountRegistryProxyConfigurables,
    /// Maximum words per blob for deployment
    pub max_words_per_blob: usize,
    /// Bytecode for the TradeAccountProxy contract
    pub trade_account_proxy_bytecode: Vec<u8>,
    /// Blob id for trading account
    pub trade_account_proxy_root: Bytes32,
    /// The storage slots of the trade account proxy
    pub trade_account_proxy_storage_slots: Vec<StorageSlot>,
    /// Bytecode for the TrialTradeAccountProxy contract
    pub trial_trade_account_proxy_bytecode: Vec<u8>,
    /// Blob id for trial trading account
    pub trial_trade_account_proxy_root: Bytes32,
    /// The storage slots of the trial trade account proxy
    pub trial_trade_account_proxy_storage_slots: Vec<StorageSlot>,
    /// Owner for the TradeAccountProxyRegistry contract
    pub proxy_owner: Option<Identity>,
    /// Owner for the TradeAccountRegistry contract
    pub registry_owner: Option<Identity>,
    /// Account allowed to register and fund trial trade accounts
    pub trial_trade_account_creator: Option<Identity>,
    /// Salt for contract deployment
    pub salt: Salt,
}

pub struct TradeAccountRegistryBlob {
    /// The ID of the deployed blob
    pub id: BlobId,
    /// Whether the blob already exists
    pub exists: bool,
    /// The blob data containing the contract bytecode
    pub blob: Blob,
}

impl Default for TradeAccountRegistryDeployConfig {
    fn default() -> Self {
        let trade_account_proxy_storage_slots: Vec<StorageSlot> =
            serde_json::from_slice(TRADE_ACCOUNT_REGISTER_PROXY_STORAGE).unwrap();
        let trade_account_proxy_bytecode = Contract::regular(
            TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            Salt::zeroed(),
            trade_account_proxy_storage_slots.clone(),
        );
        let trade_account_proxy_root = trade_account_proxy_bytecode.code_root();
        let trial_trade_account_proxy_storage_slots: Vec<StorageSlot> =
            serde_json::from_slice(TRIAL_TRADE_ACCOUNT_PROXY_STORAGE).unwrap();
        let trial_trade_account_proxy_bytecode = Contract::regular(
            TRIAL_TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            Salt::zeroed(),
            trial_trade_account_proxy_storage_slots.clone(),
        );
        let trial_trade_account_proxy_root =
            trial_trade_account_proxy_bytecode.code_root();

        Self {
            registry_bytecode: TRADE_ACCOUNT_REGISTER_BYTECODE.to_vec(),
            registry_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_REGISTER_STORAGE,
            )
            .unwrap(),
            registry_config: TradeAccountRegistryConfigurables::default(),
            registry_proxy_bytecode: TRADE_ACCOUNT_REGISTER_PROXY_BYTECODE.to_vec(),
            registry_proxy_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_REGISTER_PROXY_STORAGE,
            )
            .unwrap(),
            registry_proxy_config: TradeAccountRegistryProxyConfigurables::default(),
            max_words_per_blob: 10_000,
            trade_account_proxy_bytecode: TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            trade_account_proxy_root,
            trade_account_proxy_storage_slots,
            trial_trade_account_proxy_bytecode: TRIAL_TRADE_ACCOUNT_PROXY_BYTECODE
                .to_vec(),
            trial_trade_account_proxy_root,
            trial_trade_account_proxy_storage_slots,
            proxy_owner: None,
            registry_owner: None,
            trial_trade_account_creator: None,
            salt: Salt::default(),
        }
    }
}

/// Result of a TradeAccountRegistry deployment.
/// Contains the deployed contract instance and configuration.
#[derive(Clone)]
pub struct TradeAccountRegistryManager<W> {
    /// The TradeAccountRegistryProxy contract instance
    pub registry_proxy: TradeAccountRegistryProxy<W>,
    /// The deployed TradeAccountRegistry contract instance
    pub registry: TradeAccountRegistry<W>,
    /// Contract ID of the deployed TradeAccountRegistry
    pub contract_id: ContractId,
    /// The wallet used for deployment
    pub wallet: W,
}

impl<W> TradeAccountRegistryManager<W>
where
    W: Account + Clone,
{
    pub fn change_wallet(mut self, wallet: &W) -> Self {
        self.wallet = wallet.clone();
        self.registry = self.registry.with_account(wallet.clone());

        self
    }

    /// `base` plus the prop (margin) configurables THIS registry already
    /// carries, read back off the chain.
    ///
    /// Rebuilding the registry's implementation blob resets every
    /// configurable the deploy config does not explicitly repeat, and the
    /// five `PROP_*` values default to zero. A zeroed
    /// `PROP_ACCOUNT_ORACLE_CONTRACT_ID` makes `prop_register_contract`
    /// CALL address zero, so every margin registration reverts with
    /// `ContractNotFound` — and because only the margin SYSTEM phase
    /// writes those values, an upgrade that drops them stays broken until
    /// that phase happens to run again.
    ///
    /// Threading this through an upgrade makes it preserve whatever margin
    /// installed, whether or not the margin phase runs afterwards.
    ///
    /// A registry too old to expose `get_prop_oracle_id`, or one never
    /// margin-wired, yields `base` untouched: there is nothing to
    /// preserve, and the margin phase installs the values on its next run.
    pub async fn live_prop_config(
        &self,
        base: TradeAccountRegistryConfigurables,
    ) -> Result<TradeAccountRegistryConfigurables> {
        let Ok(oracle_id) = self
            .registry
            .methods()
            .get_prop_oracle_id()
            .simulate(Execution::state_read_only())
            .await
            .map(|result| result.value)
        else {
            return Ok(base);
        };
        if oracle_id == ContractId::zeroed() {
            return Ok(base);
        }
        let account_proxy = self
            .registry
            .methods()
            .default_prop_bytecode()
            .simulate(Execution::state_read_only())
            .await?
            .value;
        // Offsets are a property of the bundled proxy artifact, not of the
        // chain, so they are recomputed rather than read back.
        let [oracle_offset, parent_offset, index_offset] =
            crate::prop_deploy::prop_account_proxy_offsets()?;
        Ok(base
            .with_PROP_ACCOUNT_ORACLE_CONTRACT_ID(oracle_id)?
            .with_DEFAULT_PROP_ACCOUNT_PROXY(account_proxy)?
            .with_PROP_ORACLE_CONFIG_OFFSET(oracle_offset)?
            .with_PROP_PARENT_CONFIG_OFFSET(parent_offset)?
            .with_PROP_INDEX_CONFIG_OFFSET(index_offset)?)
    }

    pub fn new(deployer_wallet: W, contract_id: ContractId) -> Self {
        let proxy = TradeAccountRegistryProxy::new(contract_id, deployer_wallet.clone());
        let registry = TradeAccountRegistry::new(contract_id, deployer_wallet.clone());

        Self {
            registry_proxy: proxy,
            registry,
            contract_id,
            wallet: deployer_wallet,
        }
    }

    pub async fn register_blob(
        deployer_wallet: &W,
        trade_account_oracle_id: ContractId,
        trial_trade_account_oracle_id: ContractId,
        trade_account_proxy_blob_id: BlobId,
        trial_trade_account_proxy_blob_id: BlobId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryBlob> {
        let registry_owner = config
            .registry_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        let trial_trade_account_creator = config
            .trial_trade_account_creator
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        // Configure the contract with the trading pair assets
        let configurables = config
            .registry_config
            .clone()
            .with_ORACLE_CONTRACT_ID(trade_account_oracle_id)?
            .with_TRIAL_ORACLE_CONTRACT_ID(trial_trade_account_oracle_id)?
            .with_TRADE_ACCOUNT_PROXY_ROOT(Bits256(*config.trade_account_proxy_root))?
            .with_DEFAULT_TRADE_ACCOUNT_PROXY(ContractId::from(
                trade_account_proxy_blob_id,
            ))?
            .with_TRIAL_TRADE_ACCOUNT_PROXY_ROOT(Bits256(
                *config.trial_trade_account_proxy_root,
            ))?
            .with_DEFAULT_TRIAL_TRADE_ACCOUNT_PROXY(ContractId::from(
                trial_trade_account_proxy_blob_id,
            ))?
            .with_INITIAL_TRIAL_TRADE_ACCOUNT_CREATOR(trial_trade_account_creator)?
            .with_INITIAL_OWNER(State::Initialized(registry_owner))?;
        let blobs = Contract::regular(
            config.registry_bytecode.clone(),
            config.salt,
            config.registry_storage_slots.clone(),
        )
        .with_configurables(configurables.clone())
        .convert_to_loader(config.max_words_per_blob)?
        .blobs()
        .to_vec();
        let blob = blobs[0].clone();
        let blob_id = blob.id();
        let blob_exists = deployer_wallet.try_provider()?.blob_exists(blob_id).await?;

        Ok(TradeAccountRegistryBlob {
            id: blob_id,
            exists: blob_exists,
            blob: blob.clone(),
        })
    }

    pub async fn register_proxy_blob(
        deployer_wallet: &W,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryBlob> {
        let blobs = Contract::regular(
            config.trade_account_proxy_bytecode.clone(),
            Salt::zeroed(),
            config.trade_account_proxy_storage_slots.clone(),
        )
        .convert_to_loader(config.max_words_per_blob)?
        .blobs()
        .to_vec();
        let blob = blobs[0].clone();
        let blob_id = blob.id();
        let blob_exists = deployer_wallet.try_provider()?.blob_exists(blob_id).await?;

        Ok(TradeAccountRegistryBlob {
            id: blob_id,
            exists: blob_exists,
            blob: blob.clone(),
        })
    }

    pub async fn register_trial_proxy_blob(
        deployer_wallet: &W,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryBlob> {
        let blobs = Contract::regular(
            config.trial_trade_account_proxy_bytecode.clone(),
            Salt::zeroed(),
            config.trial_trade_account_proxy_storage_slots.clone(),
        )
        .convert_to_loader(config.max_words_per_blob)?
        .blobs()
        .to_vec();
        let blob = blobs[0].clone();
        let blob_id = blob.id();
        let blob_exists = deployer_wallet.try_provider()?.blob_exists(blob_id).await?;

        Ok(TradeAccountRegistryBlob {
            id: blob_id,
            exists: blob_exists,
            blob: blob.clone(),
        })
    }

    async fn upload_blob(
        deployer_wallet: &W,
        register_blob: TradeAccountRegistryBlob,
    ) -> Result<()> {
        if !register_blob.exists {
            let mut builder =
                BlobTransactionBuilder::default().with_blob(register_blob.blob);
            deployer_wallet.adjust_for_fee(&mut builder, 0).await?;
            deployer_wallet.add_witnesses(&mut builder)?;

            let tx = builder.build(&deployer_wallet.try_provider()?).await?;

            deployer_wallet
                .try_provider()?
                .send_transaction_and_await_commit(tx)
                .await?
                .check(None)?;
        }

        Ok(())
    }

    pub async fn deploy_register_blob(
        deployer_wallet: &W,
        trade_account_oracle_id: ContractId,
        trial_trade_account_oracle_id: ContractId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<BlobId> {
        let proxy_blob = Self::register_proxy_blob(deployer_wallet, config).await?;
        let proxy_blob_id = proxy_blob.id;
        Self::upload_blob(deployer_wallet, proxy_blob).await?;
        let trial_proxy_blob =
            Self::register_trial_proxy_blob(deployer_wallet, config).await?;
        let trial_proxy_blob_id = trial_proxy_blob.id;
        Self::upload_blob(deployer_wallet, trial_proxy_blob).await?;

        let register_blob = Self::register_blob(
            deployer_wallet,
            trade_account_oracle_id,
            trial_trade_account_oracle_id,
            proxy_blob_id,
            trial_proxy_blob_id,
            config,
        )
        .await?;
        let register_blob_id = register_blob.id;
        Self::upload_blob(deployer_wallet, register_blob).await?;

        Ok(register_blob_id)
    }

    pub async fn deploy_register_proxy(
        deployer_wallet: &W,
        registry_blob_id: &BlobId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<(TradeAccountRegistryProxy<W>, bool)> {
        let proxy_owner = config
            .proxy_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        let configurables = config
            .registry_proxy_config
            .clone()
            .with_INITIAL_OWNER(State::Initialized(proxy_owner))?
            .with_INITIAL_TARGET(ContractId::new(*registry_blob_id))?;
        let contract = Contract::regular(
            config.registry_proxy_bytecode.clone(),
            config.salt,
            config.registry_proxy_storage_slots.clone(),
        )
        .with_configurables(configurables);
        let contract_id = contract.contract_id();
        let already_deployed = deployer_wallet
            .try_provider()?
            .contract_exists(&contract_id)
            .await?;
        let proxy = TradeAccountRegistryProxy::new(contract_id, deployer_wallet.clone());

        if !already_deployed {
            contract
                .deploy(deployer_wallet, TxPolicies::default())
                .await?;
        }
        let requires_initialization = !already_deployed;

        Ok((proxy, requires_initialization))
    }

    /// Deploys a TradeAccountRegistry contract.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays gas fees)
    /// * `config` - Deployment configuration containing bytecode and settings
    /// * `oracle_id` - The oracle contract ID to configure
    ///
    /// # Returns
    /// * `Ok(TradeAccountRegistryDeploy)` - Complete deployment result with contract instance
    /// * `Err(anyhow::Error)` - If deployment or initialization fails
    ///
    /// # Process
    /// 1. Deploys the TradeAccountRegistry contract
    /// 2. Initializes the contract with the specified owner
    pub async fn deploy(
        deployer_wallet: &W,
        trade_account_oracle_id: ContractId,
        trial_trade_account_oracle_id: ContractId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryManager<W>> {
        let register_blob_id = Self::deploy_register_blob(
            deployer_wallet,
            trade_account_oracle_id,
            trial_trade_account_oracle_id,
            config,
        )
        .await?;
        let (proxy, requires_initialization) =
            TradeAccountRegistryManager::deploy_register_proxy(
                deployer_wallet,
                &register_blob_id,
                config,
            )
            .await?;
        let register_deploy = TradeAccountRegistryManager::new(
            deployer_wallet.clone(),
            proxy.contract_id(),
        );

        // If the contract is already deployed, we skip deployment
        if requires_initialization {
            // Initialize the proxy contract
            register_deploy
                .registry_proxy
                .methods()
                .initialize_proxy()
                .call()
                .await?;
            // Initialize the register with the specified owner
            register_deploy
                .registry
                .methods()
                .initialize()
                .call()
                .await?;
        }

        Ok(register_deploy)
    }

    /// Upgrades the TradeAccountRegistry contract with new bytecode.
    ///
    /// # Arguments
    /// * `config` - New deployment configuration containing updated bytecode and settings
    ///
    /// # Returns
    /// * `Ok(BlobId)` - The ID of the new blob containing the upgraded contract
    /// * `Err(anyhow::Error)` - If the upgrade fails
    ///
    /// # Process
    /// 1. Deploys new blob with updated bytecode
    /// 2. Updates the proxy to point to the new blob
    pub async fn upgrade(
        &self,
        trade_account_oracle_id: ContractId,
        trial_trade_account_oracle_id: ContractId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<BlobId> {
        // Deploy new blob with updated bytecode
        let new_blob_id = Self::deploy_register_blob(
            &self.wallet,
            trade_account_oracle_id,
            trial_trade_account_oracle_id,
            config,
        )
        .await?;

        // Update the proxy to point to the new blob
        self.registry_proxy
            .methods()
            .set_proxy_target(ContractId::new(new_blob_id))
            .call()
            .await?;

        Ok(new_blob_id)
    }

    pub async fn get_trade_account(&self, owner: Identity) -> Result<Option<ContractId>> {
        Ok(self
            .registry
            .methods()
            .get_trade_account(owner)
            .simulate(Execution::state_read_only())
            .await?
            .value)
    }

    pub async fn oracle_id(&self) -> Result<ContractId> {
        Ok(self
            .registry
            .methods()
            .get_oracle_id()
            .simulate(Execution::state_read_only())
            .await?
            .value)
    }

    pub async fn trial_oracle_id(&self) -> Result<ContractId> {
        Ok(self
            .registry
            .methods()
            .get_trial_oracle_id()
            .simulate(Execution::state_read_only())
            .await?
            .value)
    }

    pub async fn get_trial_trade_account_creator(&self) -> Result<Identity> {
        Ok(self
            .registry
            .methods()
            .get_trial_trade_account_creator()
            .call()
            .await?
            .value)
    }

    pub async fn set_trial_trade_account_creator(
        &self,
        trial_trade_account_creator: Identity,
    ) -> Result<u64> {
        Ok(self
            .registry
            .methods()
            .set_trial_trade_account_creator(trial_trade_account_creator)
            .call()
            .await?
            .tx_status
            .total_gas)
    }
}

impl TradeAccountRegistryManager<Wallet> {
    pub async fn register_trade_account(
        &self,
        owner_identity: Identity,
        trade_account_id: ContractId,
        trade_account_oracle_id: ContractId,
        call_option: &CallOption,
        deploy_config: &DeployConfig,
    ) -> Result<(ContractId, u64)> {
        let configurables: Vec<fuels::core::Configurable> = match deploy_config {
            DeployConfig::Latest(_) => TradingAccountProxyConfigurables::default()
                .with_ORACLE_CONTRACT_ID(trade_account_oracle_id)?
                .with_INITIAL_OWNER(TradingAccountState::Initialized(owner_identity))?
                .into(),
        };
        let call_handler = self
            .registry
            .methods()
            .register_contract(
                trade_account_id,
                Some(
                    configurables
                        .into_iter()
                        .map(|c| (c.offset, c.data))
                        .collect(),
                ),
            )
            .with_contract_ids(&[trade_account_id, trade_account_oracle_id]);

        let result = match call_option {
            CallOption::AwaitBlock => call_handler.call().await?,
            CallOption::AwaitPreconfirmation(ops) => {
                call_handler
                    .almost_sync_call(
                        &ops.data_builder,
                        &ops.utxo_manager,
                        &ops.tx_config,
                        None,
                        &[],
                    )
                    .await?
                    .tx_status?
            }
        };

        match result.value {
            Ok(_) => Ok((trade_account_id, result.tx_status.total_gas)),
            Err(e) => Err(anyhow::anyhow!("Failed to register trade account: {:?}", e)),
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub async fn register_trial_trade_account(
        &self,
        initial_session_id: Identity,
        trial_trade_account_id: ContractId,
        trial_trade_account_oracle_id: ContractId,
        trial_duration: u64,
        liquidation_threshold_amount: u64,
        liquidation_threshold_asset_id: AssetId,
        trial_allowed_order_book_ids: Vec<ContractId>,
        trial_user_profit_share_bps: u64,
        trial_platform_payout_identity: Identity,
        call_parameters: Option<CallParameters>,
        call_option: &CallOption,
    ) -> Result<(Bits256, u64)> {
        let configurables: Vec<fuels::core::Configurable> =
            TrialTradingAccountProxyConfigurables::default()
                .with_ORACLE_CONTRACT_ID(trial_trade_account_oracle_id)?
                .with_INITIAL_SESSION_ID(initial_session_id)?
                .with_TRIAL_TRADE_ACCOUNT_REGISTRY(self.contract_id)?
                .into();
        let mut call_handler = self
            .registry
            .methods()
            .trial_register_contract(
                trial_trade_account_id,
                Some(
                    configurables
                        .into_iter()
                        .map(|c| (c.offset, c.data))
                        .collect(),
                ),
                trial_duration,
                liquidation_threshold_amount,
                liquidation_threshold_asset_id,
                trial_allowed_order_book_ids,
                trial_user_profit_share_bps,
                trial_platform_payout_identity,
            )
            .with_contract_ids(&[trial_trade_account_id, trial_trade_account_oracle_id]);
        if let Some(call_parameters) = call_parameters {
            call_handler = call_handler.call_params(call_parameters)?;
        }

        let result = match call_option {
            CallOption::AwaitBlock => call_handler.call().await?,
            CallOption::AwaitPreconfirmation(ops) => {
                call_handler
                    .almost_sync_call(
                        &ops.data_builder,
                        &ops.utxo_manager,
                        &ops.tx_config,
                        None,
                        &[],
                    )
                    .await?
                    .tx_status?
            }
        };

        match result.value {
            Ok(root) => Ok((root, result.tx_status.total_gas)),
            Err(err) => Err(anyhow::anyhow!("Registration failed: {}", err)),
        }
    }
}

#[cfg(test)]
mod tests_trade_account_registry {
    use super::*;
    use crate::{
        trade_account::TradeAccountManager,
        trade_account_deploy::{
            DeployConfig,
            TradeAccountDeploy,
            TradeAccountDeployConfig,
        },
    };
    use fuels::test_helpers::{
        WalletsConfig,
        launch_custom_provider_and_get_wallets,
    };

    #[tokio::test]
    async fn test_trade_account_registry_upgrade() {
        let mut wallets = launch_custom_provider_and_get_wallets(
            WalletsConfig::new(Some(2), Some(1), Some(1_000_000_000)),
            None,
            Some(::fuels::test_helpers::ChainConfig::local_testnet()),
        )
        .await
        .unwrap();
        let deployer_wallet = wallets.pop().unwrap();
        let user1_wallet = wallets.pop().unwrap();

        let deploy_config = DeployConfig::Latest(TradeAccountDeployConfig::default());
        let trade_account_deploy =
            TradeAccountDeploy::deploy(&deployer_wallet, &deploy_config)
                .await
                .unwrap();
        let user1_identity = Identity::Address(user1_wallet.address());

        // Test register of trade account for user 1
        let user1_trade_account_deployed = trade_account_deploy
            .deploy_with_account(
                &user1_identity,
                &deploy_config,
                &CallOption::AwaitBlock,
                None,
                &[],
            )
            .await
            .unwrap();
        let user1_trade_account = TradeAccountManager::new(
            &user1_wallet,
            user1_trade_account_deployed.proxy_id.unwrap(),
        )
        .await
        .unwrap();

        let trade_account_registry = TradeAccountRegistryManager::deploy(
            &deployer_wallet,
            trade_account_deploy.oracle_id,
            ContractId::zeroed(),
            &Default::default(),
        )
        .await
        .unwrap();

        trade_account_registry
            .register_trade_account(
                user1_identity,
                user1_trade_account.contract_id(),
                trade_account_deploy.oracle_id,
                &CallOption::AwaitBlock,
                &deploy_config,
            )
            .await
            .unwrap();

        let result = trade_account_registry
            .get_trade_account(user1_identity)
            .await
            .unwrap();
        assert_eq!(
            result,
            Some(user1_trade_account.contract_id()),
            "Trade account should register successfully"
        );

        // Upgrate trade account regsitry
        let current_blob_id = trade_account_registry
            .registry_proxy
            .methods()
            .proxy_target()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;
        trade_account_registry
            .upgrade(
                trade_account_deploy.oracle_id,
                ContractId::zeroed(),
                &TradeAccountRegistryDeployConfig {
                    trade_account_proxy_root: Bytes32::new([1; 32]),
                    ..Default::default()
                },
            )
            .await
            .unwrap();
        let upgraded_blob_id = trade_account_registry
            .registry_proxy
            .methods()
            .proxy_target()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;
        assert_ne!(
            current_blob_id, upgraded_blob_id,
            "Blob id should be updated after upgrade"
        );

        // Test register of trade account for user 2 after upgrade
        let result = trade_account_registry
            .get_trade_account(user1_identity)
            .await
            .unwrap();
        assert_eq!(
            result,
            Some(user1_trade_account.contract_id()),
            "Trade account registered should keep working after upgrade"
        );
    }

    #[tokio::test]
    async fn test_trade_account_registration_on_registry() {
        // Start fuel-core
        let mut wallets = launch_custom_provider_and_get_wallets(
            WalletsConfig::new(Some(2), Some(1), Some(1_000_000_000)),
            None,
            Some(::fuels::test_helpers::ChainConfig::local_testnet()),
        )
        .await
        .unwrap();
        let deployer_wallet = wallets.pop().unwrap();
        let user_wallet = wallets.pop().unwrap();

        let deploy_config = DeployConfig::Latest(TradeAccountDeployConfig::default());
        let trade_account_deploy =
            TradeAccountDeploy::deploy(&deployer_wallet, &deploy_config)
                .await
                .unwrap();
        let owner_identity = Identity::Address(user_wallet.address());
        let trade_account_deployed = trade_account_deploy
            .deploy_with_account(
                &owner_identity,
                &deploy_config,
                &CallOption::AwaitBlock,
                None,
                &[],
            )
            .await
            .unwrap();
        let trade_account = TradeAccountManager::new(
            &user_wallet,
            trade_account_deployed.proxy_id.unwrap(),
        )
        .await
        .unwrap();

        let trade_account_registry = TradeAccountRegistryManager::deploy(
            &deployer_wallet,
            trade_account_deploy.oracle_id,
            ContractId::zeroed(),
            &Default::default(),
        )
        .await
        .unwrap();

        trade_account_registry
            .register_trade_account(
                owner_identity,
                trade_account.contract_id(),
                trade_account_deploy.oracle_id,
                &CallOption::AwaitBlock,
                &deploy_config,
            )
            .await
            .unwrap();

        let result = trade_account_registry
            .get_trade_account(owner_identity)
            .await
            .unwrap();
        assert_eq!(result, Some(trade_account.contract_id()));

        trade_account_registry
            .register_trade_account(
                owner_identity,
                trade_account.contract_id(),
                trade_account_deploy.oracle_id,
                &CallOption::AwaitBlock,
                &deploy_config,
            )
            .await
            .expect_err("Should fail if already register");
    }
}