// 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(); } }