smom-dbis-138/test/smart-accounts/AccountWalletRegistryExtendedTest.t.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import {Test} from "forge-std/Test.sol";
import {AccountWalletRegistryExtended} from "../../contracts/smart-accounts/AccountWalletRegistryExtended.sol";
import {IAccountWalletRegistry} from "../../contracts/emoney/interfaces/IAccountWalletRegistry.sol";

/**
 * @title AccountWalletRegistryExtendedTest
 * @notice Tests for AccountWalletRegistryExtended with Smart Account support
 */
contract AccountWalletRegistryExtendedTest is Test {
    AccountWalletRegistryExtended public registry;
    address public admin;
    address public accountManager;
    address public smartAccountFactory;
    address public entryPoint;
    
    // Mock smart account (contract address)
    address public mockSmartAccount;
    
    bytes32 public accountRefId1 = keccak256("account1");
    bytes32 public providerSmart = keccak256("METAMASK_SMART_ACCOUNT");
    bytes32 public providerEOA = keccak256("METAMASK");

    function setUp() public {
        admin = address(0x1);
        accountManager = address(0x2);
        smartAccountFactory = address(0x3);
        entryPoint = address(0x4);
        
        // Deploy mock smart account (contract)
        mockSmartAccount = address(new MockSmartAccount());
        
        registry = new AccountWalletRegistryExtended(
            admin,
            smartAccountFactory,
            entryPoint
        );

        vm.startPrank(admin);
        registry.grantRole(registry.ACCOUNT_MANAGER_ROLE(), accountManager);
        registry.grantRole(registry.ACCOUNT_MANAGER_ROLE(), address(registry)); // Self-call from linkSmartAccount->linkAccountToWallet
        vm.stopPrank();
    }

    function test_linkSmartAccount() public {
        vm.prank(accountManager);
        registry.linkSmartAccount(accountRefId1, mockSmartAccount, providerSmart);

        bytes32 walletRefId = keccak256(abi.encodePacked(mockSmartAccount));
        assertTrue(registry.isLinked(accountRefId1, walletRefId));
        assertTrue(registry.isActive(accountRefId1, walletRefId));
        assertTrue(registry.isSmartAccount(walletRefId));
        assertTrue(registry.isSmartAccountAddress(mockSmartAccount));
    }

    function test_linkSmartAccount_revertsIfNotContract() public {
        address eoa = address(0x100);
        
        vm.prank(accountManager);
        vm.expectRevert("AccountWalletRegistryExtended: not a contract");
        registry.linkSmartAccount(accountRefId1, eoa, providerSmart);
    }

    function test_linkSmartAccount_revertsIfZeroAddress() public {
        vm.prank(accountManager);
        vm.expectRevert("AccountWalletRegistryExtended: zero smartAccount");
        registry.linkSmartAccount(accountRefId1, address(0), providerSmart);
    }

    function test_isSmartAccount() public {
        bytes32 walletRefId = keccak256(abi.encodePacked(mockSmartAccount));
        
        // Initially false
        assertFalse(registry.isSmartAccount(walletRefId));
        
        // Link smart account
        vm.prank(accountManager);
        registry.linkSmartAccount(accountRefId1, mockSmartAccount, providerSmart);
        
        // Now true
        assertTrue(registry.isSmartAccount(walletRefId));
    }

    function test_isSmartAccountAddress() public {
        // Initially false
        assertFalse(registry.isSmartAccountAddress(mockSmartAccount));
        
        // Link smart account
        vm.prank(accountManager);
        registry.linkSmartAccount(accountRefId1, mockSmartAccount, providerSmart);
        
        // Now true
        assertTrue(registry.isSmartAccountAddress(mockSmartAccount));
    }

    function test_setSmartAccountFactory() public {
        address newFactory = address(0x5);
        
        vm.prank(admin);
        registry.setSmartAccountFactory(newFactory);
        
        assertEq(registry.smartAccountFactory(), newFactory);
    }

    function test_setSmartAccountFactory_revertsIfNotAdmin() public {
        address newFactory = address(0x5);
        
        vm.prank(accountManager);
        vm.expectRevert();
        registry.setSmartAccountFactory(newFactory);
    }

    function test_setEntryPoint() public {
        address newEntryPoint = address(0x6);
        
        vm.prank(admin);
        registry.setEntryPoint(newEntryPoint);
        
        assertEq(registry.entryPoint(), newEntryPoint);
    }

    function test_setEntryPoint_revertsIfNotAdmin() public {
        address newEntryPoint = address(0x6);
        
        vm.prank(accountManager);
        vm.expectRevert();
        registry.setEntryPoint(newEntryPoint);
    }

    function test_supportsBothEOAAndSmartAccount() public {
        address eoa = address(0x100);
        bytes32 eoaWalletRefId = keccak256(abi.encodePacked(eoa));
        
        // Link EOA (using parent function)
        vm.prank(accountManager);
        registry.linkAccountToWallet(accountRefId1, eoaWalletRefId, providerEOA);
        
        // Link Smart Account
        vm.prank(accountManager);
        registry.linkSmartAccount(accountRefId1, mockSmartAccount, providerSmart);
        
        // Both should be linked
        assertTrue(registry.isLinked(accountRefId1, eoaWalletRefId));
        assertTrue(registry.isLinked(accountRefId1, keccak256(abi.encodePacked(mockSmartAccount))));
        
        // Only smart account should be marked as smart account
        assertFalse(registry.isSmartAccount(eoaWalletRefId));
        assertTrue(registry.isSmartAccount(keccak256(abi.encodePacked(mockSmartAccount))));
    }

    // Events
    event SmartAccountLinked(
        bytes32 indexed accountRefId,
        address indexed smartAccount,
        bytes32 provider
    );
}

// Mock Smart Account contract for testing
contract MockSmartAccount {
    // Empty contract to simulate smart account
}