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 {
        erc721Mock.mint(USER);
		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();
	}

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

		// mint ERC721 with new TOKEN_ID
		erc721Mock.mint(USER2);
		uint256 newTokenId = TOKEN_ID + 1;

		vm.startPrank(USER2);
		erc721Mock.approve(address(tokenDivider), newTokenId);
		tokenDivider.divideNft(
			address(erc721Mock),
			newTokenId,
			AMOUNT
		);
		vm.stopPrank();

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

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

		// USER2 claim NFT
		vm.startPrank(USER2);
		erc20MockWithNewTokenId.approve(
			address(tokenDivider),
			AMOUNT
		);
		tokenDivider.claimNft(address(erc721Mock));
		vm.stopPrank();

		assertEq(
			erc721Mock.ownerOf(newTokenId),
			address(USER2)
		);
		assertEq(
			erc721Mock.ownerOf(TOKEN_ID),
			address(tokenDivider)
		);

		// assert ERC20 address before USER claim NFT
		// @audit-report: it returns different erc20 address
		// @dev: feel free to uncomment to see the error
		/*
		assertEq(
			address(
				tokenDivider.getErc20InfoFromNft(
					address(erc721Mock)
				).erc20Address
			),
			address(erc20Mock)
		);
		 */

		// USER claim NFT
		vm.startPrank(USER2);
		erc20MockWithNewTokenId.approve(
			address(tokenDivider),
			AMOUNT
		);
		// @dev: it weirdly returns ERC721InsufficientApproval
		// despite its already approved in above.
		// maybe because of different ERC20 address?
		// @dev: changing the ERC20 address and the USER prank
		// doesn't help
		/*
		tokenDivider.claimNft(address(erc721Mock));
		vm.stopPrank();

		assertEq(
			erc721Mock.ownerOf(TOKEN_ID),
			address(USER)
		);
		 */
	}
}