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

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

import {Test} 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';

contract TokenDividerClaimTest 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);

		// divide the ERC721 and mint ERC20
		vm.startPrank(USER);
		erc721Mock.approve(address(tokenDivider), TOKEN_ID);
		tokenDivider.divideNft(
			address(erc721Mock),
			TOKEN_ID,
			AMOUNT
		);
		vm.stopPrank();
    }

	function test_assertSetUp() public view {
		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_claimNftSuccessfully() public {
		ERC20Mock erc20Mock = ERC20Mock(
			tokenDivider.getErc20InfoFromNft(
				address(erc721Mock)
			).erc20Address
		);

		vm.startPrank(USER);
		erc20Mock.approve(address(tokenDivider), AMOUNT);
		tokenDivider.claimNft(address(erc721Mock));
		vm.stopPrank();

		assertEq(
			tokenDivider.getBalanceOf(
				address(USER),
				address(erc20Mock)
			),
			0			
		);
		assertEq(
			tokenDivider.getErc20TotalMintedAmount(
				address(erc20Mock)
			),
			0			
		);
		assertEq(
			erc721Mock.ownerOf(TOKEN_ID),
			address(USER)
		);
		assertEq(
			erc20Mock.balanceOf(
				address(USER)
			),	
			0
		);
	}

	function test_failIf_claimNftWithCounterfeitNFT() public {
		ERC20Mock erc20Mock = ERC20Mock(
			tokenDivider.getErc20InfoFromNft(
				address(erc721Mock)
			).erc20Address
		);

		vm.startPrank(USER2);
		erc20Mock.mint(address(USER2), AMOUNT);
		erc20Mock.approve(address(tokenDivider), AMOUNT);
		vm.expectRevert(
			TokenDivider.TokenDivider__NotEnoughErc20Balance.selector
		);
		tokenDivider.claimNft(address(erc721Mock));
		vm.stopPrank();
	}

	function test_failIf_notExistingNftAddress() public {
		ERC20Mock erc20Mock = ERC20Mock(
			tokenDivider.getErc20InfoFromNft(
				address(erc721Mock)
			).erc20Address
		);

		vm.startPrank(USER);
		erc20Mock.approve(address(tokenDivider), AMOUNT);
		tokenDivider.claimNft(address(1));
		vm.stopPrank();

		assertEq(
			tokenDivider.getBalanceOf(
				address(USER),
				address(erc20Mock)
			),
			0			
		);
		assertEq(
			tokenDivider.getErc20TotalMintedAmount(
				address(erc20Mock)
			),
			0			
		);
		assertEq(
			erc721Mock.ownerOf(TOKEN_ID),
			address(USER)
		);
		assertEq(
			erc20Mock.balanceOf(
				address(USER)
			),	
			0
		);
	}

	function test_failIf_ERC20SentToOtherOnClaim() public {
		ERC20Mock erc20Mock = ERC20Mock(
			tokenDivider.getErc20InfoFromNft(
				address(erc721Mock)
			).erc20Address
		);

		assertEq(
			erc20Mock.balanceOf(
				address(USER)
			),	
			AMOUNT
		);

		vm.startPrank(USER);
		erc20Mock.approve(address(tokenDivider), AMOUNT);
		erc20Mock.transfer(
			address(USER2),
			AMOUNT
		);
		vm.stopPrank();

		assertEq(
			tokenDivider.getBalanceOf(
				address(USER),
				address(erc20Mock)
			),
			AMOUNT			
		);
		assertEq(
			tokenDivider.getBalanceOf(
				address(USER2),
				address(erc20Mock)
			),
			0			
		);

		vm.startPrank(USER);
		erc20Mock.approve(address(tokenDivider), AMOUNT);
		vm.expectRevert();
		tokenDivider.claimNft(address(erc721Mock));
		vm.stopPrank();
	}
}