synor/crates/synor-bridge/src/vault.rs

//! Bridge Vault
//!
//! Manages locked assets for cross-chain transfers.
//! Assets are locked in vaults on the source chain and
//! wrapped tokens are minted on the destination chain.

use crate::{AssetId, BridgeAddress, BridgeError, BridgeResult, ChainType};
use borsh::{BorshDeserialize, BorshSerialize};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fmt;

/// Unique vault identifier
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct VaultId(pub String);

impl VaultId {
    /// Create a new vault ID
    pub fn new(id: impl Into<String>) -> Self {
        Self(id.into())
    }

    /// Generate vault ID from asset
    pub fn from_asset(asset: &AssetId, chain: &ChainType) -> Self {
        let mut hasher = Sha256::new();
        hasher.update(chain.to_string().as_bytes());
        hasher.update(asset.identifier.as_bytes());
        Self(hex::encode(&hasher.finalize()[..16]))
    }
}

impl fmt::Display for VaultId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Locked asset record
#[derive(Debug, Clone, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct LockedAsset {
    /// Asset being locked
    pub asset: AssetId,
    /// Amount locked
    pub amount: u128,
    /// Who locked the asset
    pub owner: BridgeAddress,
    /// Recipient on destination chain
    pub recipient: BridgeAddress,
    /// Lock timestamp (Unix seconds)
    pub locked_at: u64,
    /// Lock expiry (for timelock, 0 = no expiry)
    pub expires_at: u64,
    /// Transaction hash on source chain
    pub lock_tx_hash: Option<Vec<u8>>,
    /// Whether the asset has been released
    pub released: bool,
}

impl LockedAsset {
    /// Create a new locked asset record
    pub fn new(
        asset: AssetId,
        amount: u128,
        owner: BridgeAddress,
        recipient: BridgeAddress,
        locked_at: u64,
    ) -> Self {
        Self {
            asset,
            amount,
            owner,
            recipient,
            locked_at,
            expires_at: 0,
            lock_tx_hash: None,
            released: false,
        }
    }

    /// Set expiry time
    pub fn with_expiry(mut self, expires_at: u64) -> Self {
        self.expires_at = expires_at;
        self
    }

    /// Set lock transaction hash
    pub fn with_tx_hash(mut self, tx_hash: Vec<u8>) -> Self {
        self.lock_tx_hash = Some(tx_hash);
        self
    }

    /// Check if the lock has expired
    pub fn is_expired(&self, current_time: u64) -> bool {
        self.expires_at > 0 && current_time >= self.expires_at
    }

    /// Mark as released
    pub fn release(&mut self) {
        self.released = true;
    }
}

/// Vault state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum VaultState {
    /// Vault is active and accepting deposits
    Active,
    /// Vault is paused (no new deposits)
    Paused,
    /// Vault is deprecated (migration needed)
    Deprecated,
}

/// Asset vault for a specific chain
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Vault {
    /// Vault identifier
    pub id: VaultId,
    /// Chain this vault is on
    pub chain: ChainType,
    /// Asset managed by this vault
    pub asset: AssetId,
    /// Current vault state
    pub state: VaultState,
    /// Total locked amount
    pub total_locked: u128,
    /// Individual locked assets by lock ID
    locked_assets: HashMap<String, LockedAsset>,
    /// Vault address (contract address for EVM)
    pub vault_address: Option<BridgeAddress>,
    /// Admin addresses
    pub admins: Vec<BridgeAddress>,
    /// Daily limit (0 = unlimited)
    pub daily_limit: u128,
    /// Today's usage
    daily_usage: u128,
    /// Last reset timestamp
    last_reset: u64,
}

impl Vault {
    /// Create a new vault
    pub fn new(id: VaultId, chain: ChainType, asset: AssetId) -> Self {
        Self {
            id,
            chain,
            asset,
            state: VaultState::Active,
            total_locked: 0,
            locked_assets: HashMap::new(),
            vault_address: None,
            admins: Vec::new(),
            daily_limit: 0,
            daily_usage: 0,
            last_reset: 0,
        }
    }

    /// Set vault address
    pub fn with_address(mut self, address: BridgeAddress) -> Self {
        self.vault_address = Some(address);
        self
    }

    /// Set daily limit
    pub fn with_daily_limit(mut self, limit: u128) -> Self {
        self.daily_limit = limit;
        self
    }

    /// Add admin
    pub fn add_admin(&mut self, admin: BridgeAddress) {
        if !self.admins.contains(&admin) {
            self.admins.push(admin);
        }
    }

    /// Lock assets in the vault
    pub fn lock(
        &mut self,
        lock_id: impl Into<String>,
        amount: u128,
        owner: BridgeAddress,
        recipient: BridgeAddress,
        current_time: u64,
    ) -> BridgeResult<()> {
        // Check vault state
        if self.state != VaultState::Active {
            return Err(BridgeError::BridgePaused);
        }

        // Check daily limit
        self.check_daily_limit(amount, current_time)?;

        let lock_id = lock_id.into();
        if self.locked_assets.contains_key(&lock_id) {
            return Err(BridgeError::TransferAlreadyExists(lock_id));
        }

        let locked = LockedAsset::new(
            self.asset.clone(),
            amount,
            owner,
            recipient,
            current_time,
        );

        self.locked_assets.insert(lock_id, locked);
        self.total_locked += amount;
        self.daily_usage += amount;

        Ok(())
    }

    /// Unlock assets from the vault
    pub fn unlock(&mut self, lock_id: &str) -> BridgeResult<LockedAsset> {
        let locked = self
            .locked_assets
            .get_mut(lock_id)
            .ok_or_else(|| BridgeError::TransferNotFound(lock_id.to_string()))?;

        if locked.released {
            return Err(BridgeError::TransferAlreadyCompleted(lock_id.to_string()));
        }

        locked.release();
        self.total_locked = self.total_locked.saturating_sub(locked.amount);

        Ok(locked.clone())
    }

    /// Get locked asset
    pub fn get_locked(&self, lock_id: &str) -> Option<&LockedAsset> {
        self.locked_assets.get(lock_id)
    }

    /// Check and update daily limit
    fn check_daily_limit(&mut self, amount: u128, current_time: u64) -> BridgeResult<()> {
        if self.daily_limit == 0 {
            return Ok(());
        }

        // Reset daily usage if new day
        let day = current_time / 86400;
        let last_day = self.last_reset / 86400;
        if day > last_day {
            self.daily_usage = 0;
            self.last_reset = current_time;
        }

        // Check limit
        if self.daily_usage + amount > self.daily_limit {
            return Err(BridgeError::RateLimitExceeded);
        }

        Ok(())
    }

    /// Pause the vault
    pub fn pause(&mut self) {
        self.state = VaultState::Paused;
    }

    /// Resume the vault
    pub fn resume(&mut self) {
        self.state = VaultState::Active;
    }

    /// Deprecate the vault
    pub fn deprecate(&mut self) {
        self.state = VaultState::Deprecated;
    }

    /// Get all locked assets
    pub fn all_locked(&self) -> impl Iterator<Item = (&String, &LockedAsset)> {
        self.locked_assets.iter()
    }

    /// Get active (unreleased) locked assets
    pub fn active_locked(&self) -> impl Iterator<Item = (&String, &LockedAsset)> {
        self.locked_assets.iter().filter(|(_, l)| !l.released)
    }

    /// Get expired locked assets
    pub fn expired_locked(&self, current_time: u64) -> impl Iterator<Item = (&String, &LockedAsset)> {
        self.locked_assets
            .iter()
            .filter(move |(_, l)| !l.released && l.is_expired(current_time))
    }
}

/// Vault manager for multiple vaults
pub struct VaultManager {
    /// Vaults by ID
    vaults: HashMap<VaultId, Vault>,
    /// Vault lookup by (chain, asset)
    by_chain_asset: HashMap<(ChainType, String), VaultId>,
}

impl VaultManager {
    /// Create a new vault manager
    pub fn new() -> Self {
        Self {
            vaults: HashMap::new(),
            by_chain_asset: HashMap::new(),
        }
    }

    /// Create and register a new vault
    pub fn create_vault(&mut self, chain: ChainType, asset: AssetId) -> VaultId {
        let vault_id = VaultId::from_asset(&asset, &chain);

        let vault = Vault::new(vault_id.clone(), chain.clone(), asset.clone());

        self.by_chain_asset
            .insert((chain, asset.identifier.clone()), vault_id.clone());
        self.vaults.insert(vault_id.clone(), vault);

        vault_id
    }

    /// Get vault by ID
    pub fn get_vault(&self, vault_id: &VaultId) -> Option<&Vault> {
        self.vaults.get(vault_id)
    }

    /// Get mutable vault by ID
    pub fn get_vault_mut(&mut self, vault_id: &VaultId) -> Option<&mut Vault> {
        self.vaults.get_mut(vault_id)
    }

    /// Find vault by chain and asset
    pub fn find_vault(&self, chain: &ChainType, asset: &AssetId) -> Option<&Vault> {
        self.by_chain_asset
            .get(&(chain.clone(), asset.identifier.clone()))
            .and_then(|id| self.vaults.get(id))
    }

    /// Find mutable vault by chain and asset
    pub fn find_vault_mut(&mut self, chain: &ChainType, asset: &AssetId) -> Option<&mut Vault> {
        let vault_id = self
            .by_chain_asset
            .get(&(chain.clone(), asset.identifier.clone()))?
            .clone();
        self.vaults.get_mut(&vault_id)
    }

    /// Get or create vault for asset
    pub fn get_or_create_vault(&mut self, chain: ChainType, asset: AssetId) -> &mut Vault {
        let key = (chain.clone(), asset.identifier.clone());
        if !self.by_chain_asset.contains_key(&key) {
            self.create_vault(chain.clone(), asset.clone());
        }
        let vault_id = self.by_chain_asset.get(&key).unwrap().clone();
        self.vaults.get_mut(&vault_id).unwrap()
    }

    /// Get total locked value across all vaults
    pub fn total_locked(&self) -> u128 {
        self.vaults.values().map(|v| v.total_locked).sum()
    }

    /// Get total locked value for an asset
    pub fn total_locked_for_asset(&self, asset: &AssetId) -> u128 {
        self.vaults
            .values()
            .filter(|v| v.asset.identifier == asset.identifier)
            .map(|v| v.total_locked)
            .sum()
    }

    /// List all vault IDs
    pub fn vault_ids(&self) -> Vec<VaultId> {
        self.vaults.keys().cloned().collect()
    }
}

impl Default for VaultManager {
    fn default() -> Self {
        Self::new()
    }
}

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

    // ==================== Helper Functions ====================

    fn test_owner() -> BridgeAddress {
        BridgeAddress::from_eth([0xaa; 20])
    }

    fn test_recipient() -> BridgeAddress {
        BridgeAddress::from_synor([0xbb; 32])
    }

    fn test_owner_alt() -> BridgeAddress {
        BridgeAddress::from_eth([0xcc; 20])
    }

    fn test_erc20_asset() -> AssetId {
        AssetId::erc20("0x1234567890123456789012345678901234567890", "USDC", 6)
    }

    // ==================== VaultId Tests ====================

    #[test]
    fn test_vault_id_new() {
        let id = VaultId::new("my-vault");
        assert_eq!(id.0, "my-vault");
    }

    #[test]
    fn test_vault_id_display() {
        let id = VaultId::new("vault-123");
        assert_eq!(format!("{}", id), "vault-123");
    }

    #[test]
    fn test_vault_id_from_asset_deterministic() {
        let asset = AssetId::eth();
        let id1 = VaultId::from_asset(&asset, &ChainType::Ethereum);
        let id2 = VaultId::from_asset(&asset, &ChainType::Ethereum);

        assert_eq!(id1, id2);
        assert!(!id1.0.is_empty());
    }

    #[test]
    fn test_vault_id_from_asset_different_chains() {
        let asset = AssetId::eth();
        let id1 = VaultId::from_asset(&asset, &ChainType::Ethereum);
        let id2 = VaultId::from_asset(&asset, &ChainType::EthereumSepolia);

        assert_ne!(id1, id2);
    }

    #[test]
    fn test_vault_id_from_asset_different_assets() {
        let eth = AssetId::eth();
        let usdc = test_erc20_asset();

        let id1 = VaultId::from_asset(&eth, &ChainType::Ethereum);
        let id2 = VaultId::from_asset(&usdc, &ChainType::Ethereum);

        assert_ne!(id1, id2);
    }

    // ==================== LockedAsset Tests ====================

    #[test]
    fn test_locked_asset_new() {
        let locked = LockedAsset::new(
            AssetId::eth(),
            1000,
            test_owner(),
            test_recipient(),
            1700000000,
        );

        assert_eq!(locked.amount, 1000);
        assert_eq!(locked.locked_at, 1700000000);
        assert_eq!(locked.expires_at, 0);
        assert!(locked.lock_tx_hash.is_none());
        assert!(!locked.released);
    }

    #[test]
    fn test_locked_asset_with_expiry() {
        let locked = LockedAsset::new(
            AssetId::eth(),
            1000,
            test_owner(),
            test_recipient(),
            1700000000,
        )
        .with_expiry(1700003600);

        assert_eq!(locked.expires_at, 1700003600);
    }

    #[test]
    fn test_locked_asset_with_tx_hash() {
        let tx_hash = vec![0x11; 32];
        let locked = LockedAsset::new(
            AssetId::eth(),
            1000,
            test_owner(),
            test_recipient(),
            1700000000,
        )
        .with_tx_hash(tx_hash.clone());

        assert_eq!(locked.lock_tx_hash, Some(tx_hash));
    }

    #[test]
    fn test_locked_asset_is_expired_no_expiry() {
        let locked = LockedAsset::new(
            AssetId::eth(),
            1000,
            test_owner(),
            test_recipient(),
            1700000000,
        );

        assert!(!locked.is_expired(1800000000));
    }

    #[test]
    fn test_locked_asset_expiry() {
        let locked = LockedAsset::new(
            AssetId::eth(),
            1000,
            test_owner(),
            test_recipient(),
            1000,
        )
        .with_expiry(2000);

        assert!(!locked.is_expired(1500));
        assert!(locked.is_expired(2000));
        assert!(locked.is_expired(3000));
    }

    #[test]
    fn test_locked_asset_release() {
        let mut locked = LockedAsset::new(
            AssetId::eth(),
            1000,
            test_owner(),
            test_recipient(),
            1700000000,
        );

        assert!(!locked.released);
        locked.release();
        assert!(locked.released);
    }

    // ==================== VaultState Tests ====================

    #[test]
    fn test_vault_state_equality() {
        assert_eq!(VaultState::Active, VaultState::Active);
        assert_eq!(VaultState::Paused, VaultState::Paused);
        assert_eq!(VaultState::Deprecated, VaultState::Deprecated);
        assert_ne!(VaultState::Active, VaultState::Paused);
    }

    // ==================== Vault Tests ====================

    #[test]
    fn test_vault_new() {
        let vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        assert_eq!(vault.id.0, "test-vault");
        assert_eq!(vault.chain, ChainType::Ethereum);
        assert_eq!(vault.state, VaultState::Active);
        assert_eq!(vault.total_locked, 0);
        assert_eq!(vault.daily_limit, 0);
    }

    #[test]
    fn test_vault_with_address() {
        let address = test_owner();
        let vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        )
        .with_address(address.clone());

        assert_eq!(vault.vault_address, Some(address));
    }

    #[test]
    fn test_vault_with_daily_limit() {
        let vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        )
        .with_daily_limit(1000000);

        assert_eq!(vault.daily_limit, 1000000);
    }

    #[test]
    fn test_vault_add_admin() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        let admin = test_owner();
        vault.add_admin(admin.clone());

        assert_eq!(vault.admins.len(), 1);
        assert_eq!(vault.admins[0], admin);
    }

    #[test]
    fn test_vault_add_admin_duplicate() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        let admin = test_owner();
        vault.add_admin(admin.clone());
        vault.add_admin(admin.clone());

        assert_eq!(vault.admins.len(), 1);
    }

    #[test]
    fn test_lock_unlock() {
        let mut vault = Vault::new(
            VaultId::new("test"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        let current_time = 1700000000;

        vault
            .lock("lock1", 1000, test_owner(), test_recipient(), current_time)
            .unwrap();

        assert_eq!(vault.total_locked, 1000);
        assert!(vault.get_locked("lock1").is_some());

        let released = vault.unlock("lock1").unwrap();
        assert_eq!(released.amount, 1000);
        assert!(released.released);
        assert_eq!(vault.total_locked, 0);
    }

    #[test]
    fn test_vault_lock_multiple() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        let current_time = 1700000000;
        vault.lock("lock-1", 1000, test_owner(), test_recipient(), current_time).unwrap();
        vault.lock("lock-2", 2000, test_owner(), test_recipient(), current_time).unwrap();
        vault.lock("lock-3", 500, test_owner_alt(), test_recipient(), current_time).unwrap();

        assert_eq!(vault.total_locked, 3500);
    }

    #[test]
    fn test_duplicate_lock() {
        let mut vault = Vault::new(
            VaultId::new("test"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault
            .lock("lock1", 1000, test_owner(), test_recipient(), 0)
            .unwrap();

        let result = vault.lock("lock1", 500, test_owner(), test_recipient(), 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_vault_unlock_nonexistent() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        let result = vault.unlock("nonexistent");
        assert!(matches!(result, Err(BridgeError::TransferNotFound(_))));
    }

    #[test]
    fn test_vault_unlock_already_released() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();
        vault.unlock("lock-1").unwrap();

        let result = vault.unlock("lock-1");
        assert!(matches!(result, Err(BridgeError::TransferAlreadyCompleted(_))));
    }

    #[test]
    fn test_vault_pause() {
        let mut vault = Vault::new(
            VaultId::new("test"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.pause();

        let result = vault.lock("lock1", 1000, test_owner(), test_recipient(), 0);
        assert!(matches!(result, Err(BridgeError::BridgePaused)));
    }

    #[test]
    fn test_vault_resume() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.pause();
        vault.resume();

        assert_eq!(vault.state, VaultState::Active);
        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();
    }

    #[test]
    fn test_vault_deprecate() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.deprecate();
        assert_eq!(vault.state, VaultState::Deprecated);

        let result = vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0);
        assert!(matches!(result, Err(BridgeError::BridgePaused)));
    }

    #[test]
    fn test_daily_limit() {
        let mut vault = Vault::new(
            VaultId::new("test"),
            ChainType::Ethereum,
            AssetId::eth(),
        )
        .with_daily_limit(1000);

        let current_time = 86400 * 100;

        vault
            .lock("lock1", 500, test_owner(), test_recipient(), current_time)
            .unwrap();

        let result = vault.lock("lock2", 600, test_owner(), test_recipient(), current_time);
        assert!(matches!(result, Err(BridgeError::RateLimitExceeded)));

        let next_day = current_time + 86400;
        vault
            .lock("lock2", 600, test_owner(), test_recipient(), next_day)
            .unwrap();
    }

    #[test]
    fn test_vault_no_daily_limit() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        let current_time = 0;
        vault.lock("lock-1", 1000000000, test_owner(), test_recipient(), current_time).unwrap();
        vault.lock("lock-2", 1000000000, test_owner(), test_recipient(), current_time).unwrap();

        assert_eq!(vault.total_locked, 2000000000);
    }

    #[test]
    fn test_vault_get_locked() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();

        assert!(vault.get_locked("lock-1").is_some());
        assert!(vault.get_locked("nonexistent").is_none());
    }

    #[test]
    fn test_vault_all_locked() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();
        vault.lock("lock-2", 2000, test_owner(), test_recipient(), 0).unwrap();

        let all: Vec<_> = vault.all_locked().collect();
        assert_eq!(all.len(), 2);
    }

    #[test]
    fn test_vault_active_locked() {
        let mut vault = Vault::new(
            VaultId::new("test-vault"),
            ChainType::Ethereum,
            AssetId::eth(),
        );

        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();
        vault.lock("lock-2", 2000, test_owner(), test_recipient(), 0).unwrap();
        vault.unlock("lock-1").unwrap();

        let active: Vec<_> = vault.active_locked().collect();
        assert_eq!(active.len(), 1);
    }

    // ==================== VaultManager Tests ====================

    #[test]
    fn test_vault_manager_new() {
        let manager = VaultManager::new();
        assert_eq!(manager.total_locked(), 0);
        assert!(manager.vault_ids().is_empty());
    }

    #[test]
    fn test_vault_manager_default() {
        let manager = VaultManager::default();
        assert_eq!(manager.total_locked(), 0);
    }

    #[test]
    fn test_vault_manager() {
        let mut manager = VaultManager::new();

        let eth = AssetId::eth();
        let vault_id = manager.create_vault(ChainType::Ethereum, eth.clone());

        assert!(manager.get_vault(&vault_id).is_some());
        assert!(manager.find_vault(&ChainType::Ethereum, &eth).is_some());

        let vault = manager.get_or_create_vault(ChainType::Ethereum, eth.clone());
        vault.lock("lock1", 100, test_owner(), test_recipient(), 0).unwrap();

        assert_eq!(manager.total_locked(), 100);
    }

    #[test]
    fn test_vault_manager_create_multiple() {
        let mut manager = VaultManager::new();

        manager.create_vault(ChainType::Ethereum, AssetId::eth());
        manager.create_vault(ChainType::Ethereum, test_erc20_asset());
        manager.create_vault(ChainType::EthereumSepolia, AssetId::eth());

        assert_eq!(manager.vault_ids().len(), 3);
    }

    #[test]
    fn test_vault_manager_get_vault_mut() {
        let mut manager = VaultManager::new();
        let vault_id = manager.create_vault(ChainType::Ethereum, AssetId::eth());

        {
            let vault = manager.get_vault_mut(&vault_id).unwrap();
            vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();
        }

        let vault = manager.get_vault(&vault_id).unwrap();
        assert_eq!(vault.total_locked, 1000);
    }

    #[test]
    fn test_vault_manager_find_vault_not_found() {
        let manager = VaultManager::new();
        let vault = manager.find_vault(&ChainType::Ethereum, &AssetId::eth());
        assert!(vault.is_none());
    }

    #[test]
    fn test_vault_manager_find_vault_mut() {
        let mut manager = VaultManager::new();
        let eth = AssetId::eth();
        manager.create_vault(ChainType::Ethereum, eth.clone());

        let vault = manager.find_vault_mut(&ChainType::Ethereum, &eth).unwrap();
        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();

        assert_eq!(manager.total_locked(), 1000);
    }

    #[test]
    fn test_vault_manager_get_or_create_new() {
        let mut manager = VaultManager::new();
        let eth = AssetId::eth();

        let vault = manager.get_or_create_vault(ChainType::Ethereum, eth.clone());
        vault.lock("lock-1", 1000, test_owner(), test_recipient(), 0).unwrap();

        assert_eq!(manager.vault_ids().len(), 1);
        assert_eq!(manager.total_locked(), 1000);
    }

    #[test]
    fn test_vault_manager_total_locked() {
        let mut manager = VaultManager::new();

        let eth = AssetId::eth();
        let usdc = test_erc20_asset();

        let eth_vault_id = manager.create_vault(ChainType::Ethereum, eth);
        let usdc_vault_id = manager.create_vault(ChainType::Ethereum, usdc);

        manager
            .get_vault_mut(&eth_vault_id)
            .unwrap()
            .lock("lock-1", 1000, test_owner(), test_recipient(), 0)
            .unwrap();

        manager
            .get_vault_mut(&usdc_vault_id)
            .unwrap()
            .lock("lock-2", 2000, test_owner(), test_recipient(), 0)
            .unwrap();

        assert_eq!(manager.total_locked(), 3000);
    }

    #[test]
    fn test_vault_manager_total_locked_for_asset() {
        let mut manager = VaultManager::new();

        let eth = AssetId::eth();
        let usdc = test_erc20_asset();

        let eth_vault_id = manager.create_vault(ChainType::Ethereum, eth.clone());
        let usdc_vault_id = manager.create_vault(ChainType::Ethereum, usdc.clone());

        manager
            .get_vault_mut(&eth_vault_id)
            .unwrap()
            .lock("lock-1", 1000, test_owner(), test_recipient(), 0)
            .unwrap();

        manager
            .get_vault_mut(&usdc_vault_id)
            .unwrap()
            .lock("lock-2", 2000, test_owner(), test_recipient(), 0)
            .unwrap();

        assert_eq!(manager.total_locked_for_asset(&eth), 1000);
        assert_eq!(manager.total_locked_for_asset(&usdc), 2000);
    }

    #[test]
    fn test_vault_manager_vault_ids() {
        let mut manager = VaultManager::new();

        let id1 = manager.create_vault(ChainType::Ethereum, AssetId::eth());
        let id2 = manager.create_vault(ChainType::Ethereum, test_erc20_asset());

        let ids = manager.vault_ids();
        assert_eq!(ids.len(), 2);
        assert!(ids.contains(&id1));
        assert!(ids.contains(&id2));
    }
}