2025-01-cyfrin-pieces-protocol/test/unit/TokenDividerTest.audit.divideNft.t.sol

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


import {Test, console} from 'forge-std/Test.sol';
import {DeployTokenDivider} from 'script/DeployTokenDivider.s.sol';
import {TokenDivider} from 'src/TokenDivider.sol';
import {ERC721Mock} from '../mocks/ERC721Mock.sol';
import {ERC20Mock} from '@openzeppelin/contracts/mocks/token/ERC20Mock.sol';
import {Ownable} from '@openzeppelin/contracts/access/Ownable.sol';


contract TokenDiverTest is Test {
    DeployTokenDivider deployer;
    TokenDivider tokenDivider;
    ERC721Mock erc721Mock;

    address public USER = makeAddr("user");
    address public USER2 = makeAddr("user2");
    uint256 constant public STARTING_USER_BALANCE = 10e18;
    uint256 constant public AMOUNT = 2e18;
    uint256 constant public TOKEN_ID = 0;

    function setUp() public {
        deployer = new DeployTokenDivider();
        tokenDivider = deployer.run();
        
        erc721Mock = new ERC721Mock();
        
        erc721Mock.mint(USER); // @auditor: user has the NFT
        vm.deal(USER2, STARTING_USER_BALANCE); // @auditor: user2 got 10 ether
    }

    function test_assertSetUp() public view {
	// USER has the ERC721 NFT of TOKEN_ID
	assertTrue(address(erc721Mock.ownerOf(TOKEN_ID)) == address(USER));

	// USER2 has no ERC721 NFT of TOKEN_ID
	assertTrue(address(erc721Mock.ownerOf(TOKEN_ID)) != address(USER2));
    }

    function test_divideNftSuccessfully() public {
    	vm.startPrank(USER);

	// approve the NFT to receive the NFT in USER
        erc721Mock.approve(address(tokenDivider), TOKEN_ID);

	tokenDivider.divideNft(
		address(erc721Mock),
		TOKEN_ID,
		AMOUNT
	);
	vm.stopPrank();

        ERC20Mock erc20Mock = ERC20Mock(
		tokenDivider.getErc20InfoFromNft(address(erc721Mock)).erc20Address
	);

	// assert if NFT already moved to tokenDivider
        assertEq(
		erc721Mock.ownerOf(TOKEN_ID),
		address(tokenDivider)
	);
	// assert if USER has the ERC20 token
        assertEq(
		tokenDivider.getBalanceOf(USER,  address(erc20Mock)),
		AMOUNT
	);
    }

    function test_failIf_sendNFTDuringDivide() public {
    	vm.startPrank(USER);

	// approve the NFT to receive the NFT in USER
        erc721Mock.approve(address(tokenDivider), TOKEN_ID);

	// send the NFT to USER2
	erc721Mock.safeTransferFrom(
		address(USER),
		address(USER2),
		TOKEN_ID
	);

	vm.expectRevert(
		TokenDivider.TokenDivider__NotFromNftOwner.selector
	);
	tokenDivider.divideNft(
		address(erc721Mock),
		TOKEN_ID,
		AMOUNT
	);

	vm.stopPrank();
    }

    function test_sendNFTBeforeDivide() public {
    	vm.startPrank(USER);

	// send the NFT to USER2
	erc721Mock.safeTransferFrom(
		address(USER),
		address(USER2),
		TOKEN_ID
	);

	vm.stopPrank();

	vm.startPrank(USER2);

	// approve the NFT to receive the NFT in USER
        erc721Mock.approve(address(tokenDivider), TOKEN_ID);

	tokenDivider.divideNft(
		address(erc721Mock),
		TOKEN_ID,
		AMOUNT
	);

	vm.stopPrank();

	ERC20Mock erc20Mock = ERC20Mock(
		tokenDivider.getErc20InfoFromNft(address(erc721Mock)).erc20Address
	);

	// assert if USER2 has the ERC20 token
        assertEq(
		tokenDivider.getBalanceOf(USER2,  address(erc20Mock)),
		AMOUNT
	);
    }
    
    function test_failIf_nonExistenceNftTokenId() public {
    	vm.startPrank(USER);

	// approve the NFT to receive the NFT in USER
        erc721Mock.approve(address(tokenDivider), TOKEN_ID);

	vm.expectRevert();
	tokenDivider.divideNft(
		address(erc721Mock),
		404,
		AMOUNT
	);
	vm.stopPrank();
    }

    function test_failWhen_reMintTheDividedERC20Token() public {
    	vm.startPrank(USER);

	// approve the NFT to receive the NFT in USER
        erc721Mock.approve(address(tokenDivider), TOKEN_ID);

	tokenDivider.divideNft(
		address(erc721Mock),
		TOKEN_ID,
		AMOUNT
	);
	vm.stopPrank();

        ERC20Mock erc20Mock = ERC20Mock(
		tokenDivider.getErc20InfoFromNft(address(erc721Mock)).erc20Address
	);

	// assert if USER has the ERC20 token
        assertEq(
		tokenDivider.getBalanceOf(USER,  address(erc20Mock)),
		AMOUNT
	);

	vm.startPrank(USER2);
	erc20Mock.mint(address(USER2), AMOUNT);
	vm.stopPrank();

	assertEq(
		tokenDivider.getBalanceOf(
			USER, 
			address(erc20Mock)
		),
		AMOUNT
	);
	assertEq(
		tokenDivider.getBalanceOf(
			USER2, 
			address(erc20Mock)
		),
		erc20Mock.balanceOf(address(USER2))
	);
	assertEq(
		erc20Mock.totalSupply(),
		AMOUNT
	);
    }

    function test_divideSameNFTAddressWithDifferentTokenId() public {
    	vm.startPrank(USER);

	// approve the NFT to receive the NFT in USER
        erc721Mock.approve(address(tokenDivider), TOKEN_ID);

	tokenDivider.divideNft(
		address(erc721Mock),
		TOKEN_ID,
		AMOUNT
	);
	vm.stopPrank();

        ERC20Mock erc20Mock = ERC20Mock(
		tokenDivider.getErc20InfoFromNft(address(erc721Mock)).erc20Address
	);

	// assert if NFT already moved to tokenDivider
        assertEq(
		erc721Mock.ownerOf(TOKEN_ID),
		address(tokenDivider)
	);
	// assert if USER has the ERC20 token
        assertEq(
		tokenDivider.getBalanceOf(USER,  address(erc20Mock)),
		AMOUNT
	);

	uint256 counterfeitTokenId = TOKEN_ID + 1;
	erc721Mock.mint(USER2);
	assertEq(
		erc721Mock.ownerOf(counterfeitTokenId),
		address(USER2)
	);

	// USER2 divide the counterfeit
	vm.startPrank(USER2);

	// approve the NFT to receive the NFT in USER2
        erc721Mock.approve(
		address(tokenDivider),
		counterfeitTokenId
	);

	tokenDivider.divideNft(
		address(erc721Mock),
		counterfeitTokenId,
		AMOUNT
	);
	vm.stopPrank();

	// now both of users (USER and USER2) has the same amount of registered ERC20 tokens despite it refer to different NFT Token Id

	ERC20Mock erc20MockCounterfeit = ERC20Mock(
		tokenDivider.getErc20InfoFromNft(address(erc721Mock)).erc20Address
	);
	// assert if NFT already moved to tokenDivider

        assertEq(
		erc721Mock.ownerOf(counterfeitTokenId),
		address(tokenDivider)
	);

	// assert if previous NFT is still being hold by contract
	assertEq(
		erc721Mock.ownerOf(TOKEN_ID),
		address(tokenDivider)
	);

	// assert if USER2 has the ERC20 token
        assertEq(
		tokenDivider.getBalanceOf(
			USER2, 
			address(erc20MockCounterfeit)
	),
		AMOUNT
	);
	// assert if USER has the ERC20 token
        assertEq(
		tokenDivider.getBalanceOf(
			USER, 
			address(erc20MockCounterfeit)
	),
		AMOUNT
	);
    }

}