uniswap-v2-contracts/test/UniswapV2Router.t.sol

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

import {Test} from "forge-std/Test.sol";

import "../src/mock/MockERC20.sol";
import "../src/interfaces/IERC20.sol";

import "../src/UniswapV2Factory.sol";
import "../src/UniswapV2Router.sol";
import "../src/UniswapV2Pair.sol";
import "../src/WETH9.sol";

import "../src/libraries/UniswapV2Library.sol";

contract TestUniswapV2Router is Test {
    UniswapV2Factory public factory;
    UniswapV2Router public router;

    MockERC20 public token0;
    MockERC20 public token1;

    WETH9 public weth;

    fallback() external payable {}
    receive() external payable {}

    function setUp() public {
        weth = new WETH9();
        factory = new UniswapV2Factory(address(0));
        router = new UniswapV2Router(address(factory), address(weth));

        token0 = new MockERC20("SomeToken0", "ST0", 18);
        token1 = new MockERC20("SomeToken1", "ST1", 9);

        token0.mint(address(this), 10 ether);
        token1.mint(address(this), 10 ether);
    }

    function testAddLiquidityEthPairFor() public {
        token0.approve(address(router), 10 ether);

        (address _token0, address _token1) = UniswapV2Library.sortTokens(
            address(token0),
            address(weth)
        );
        address pair = UniswapV2Library.pairFor(
            address(factory),
            _token0,
            _token1
        );

        (, , uint256 liquidity) = router.addLiquidityETH{value: 1 ether}(
            address(token0),
            1 ether,
            1 ether,
            1 ether,
            address(this),
            block.timestamp + 1
        );

        assertEq(liquidity, 1 ether - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
        assertEq(factory.getPair(address(token0), address(weth)), pair);
        assertEq(weth.balanceOf(address(pair)), 1 ether);
        assertEq(token0.balanceOf(address(pair)), 1 ether);
    }

    function testAddLiquidityPairFor() public {
        token0.approve(address(router), 10 ether);
        token1.approve(address(router), 10 ether);

        (address _token0, address _token1) = UniswapV2Library.sortTokens(
            address(token0),
            address(token1)
        );
        address pair = UniswapV2Library.pairFor(
            address(factory),
            _token0,
            _token1
        );

        (, , uint256 liquidity) = router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            1 ether,
            1 ether,
            1 ether,
            address(this),
            block.timestamp + 1
        );

        assertEq(liquidity, 1 ether - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
        assertEq(factory.getPair(address(token0), address(token1)), pair);
    }

    function testAddLiquidityNoPair() public {
        token0.approve(address(router), 10 ether);
        token1.approve(address(router), 10 ether);

        (address _token0, address _token1) = UniswapV2Library.sortTokens(
            address(token0),
            address(token1)
        );

        address pair = UniswapV2Library.pairFor(
            address(factory),
            _token0,
            _token1
        );

        (uint256 amount0, uint256 amount1, uint256 liquidity) = router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            1 ether,
            1 ether,
            1 ether,
            address(this),
            block.timestamp + 1
        );

        assertEq(amount0, 1 ether);
        assertEq(amount1, 1 ether);
        assertEq(liquidity, 1 ether - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());

        assertEq(factory.getPair(address(token0), address(token1)), pair);
        assertEq(UniswapV2Pair(pair).token0(), address(token0));
        assertEq(UniswapV2Pair(pair).token1(), address(token1));

        (uint256 reserve0, uint256 reserve1, ) = UniswapV2Pair(pair).getReserves();
        assertEq(reserve0, 1 ether);
        assertEq(reserve1, 1 ether);
        assertEq(token0.balanceOf(address(pair)), 1 ether);
        assertEq(token1.balanceOf(address(pair)), 1 ether);
        assertEq(token0.balanceOf(address(this)), 9 ether);
        assertEq(token1.balanceOf(address(this)), 9 ether);
    }

    function testAddLiquidityInsufficientAmountB() public {
        token0.approve(address(router), 4 ether);
        token1.approve(address(router), 8 ether);

        router.addLiquidity(
            address(token0),
            address(token1),
            4 ether,
            8 ether,
            4 ether,
            8 ether,
            address(this),
            block.timestamp + 1
        );

        token0.approve(address(router), 1 ether);
        token1.approve(address(router), 2 ether);

        vm.expectRevert();
        router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            2 ether,
            1 ether,
            2.3 ether,
            address(this),
            block.timestamp + 1
        );
    }

    function testAddLiquidityAmountBDesiredHigh() public {
        token0.approve(address(router), 4 ether);
        token1.approve(address(router), 8 ether);

        router.addLiquidity(
            address(token0),
            address(token1),
            4 ether,
            8 ether,
            4 ether,
            8 ether,
            address(this),
            block.timestamp + 1
        );

        token0.approve(address(router), 1 ether);
        token1.approve(address(router), 2 ether);

        (uint256 amount0, uint256 amount1, ) = router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            2.3 ether,
            1 ether,
            2 ether,
            address(this),
            block.timestamp + 1
        );

        assertEq(amount0, 1 ether);
        assertEq(amount1, 2 ether);
    }

    function testAddLiquidityAmountBDesiredLow() public {
        token0.approve(address(router), 4 ether);
        token1.approve(address(router), 8 ether);

        router.addLiquidity(
            address(token0),
            address(token1),
            4 ether,
            8 ether,
            4 ether,
            8 ether,
            address(this),
            block.timestamp + 1
        );

        token0.approve(address(router), 1 ether);
        token1.approve(address(router), 2 ether);

        (uint256 amount0, uint256 amount1, ) = router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            1.5 ether,
            0.75 ether,
            2 ether,
            address(this),
            block.timestamp + 1
        );

        assertEq(amount0, 0.75 ether);
        assertEq(amount1, 1.5 ether);
    }

    function testAddLiquidityInsufficientAmountA() public {
        token0.approve(address(router), 4 ether);
        token1.approve(address(router), 8 ether);

        router.addLiquidity(
            address(token0),
            address(token1),
            4 ether,
            8 ether,
            4 ether,
            8 ether,
            address(this),
            block.timestamp + 1
        );

        token0.approve(address(router), 1 ether);
        token1.approve(address(router), 2 ether);

        vm.expectRevert();
        router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            1.5 ether,
            1 ether,
            2 ether,
            address(this),
            block.timestamp + 1
        );
    }

    function testAddLiquidityExpired() public {
        token0.approve(address(router), 1 ether);
        token1.approve(address(router), 1 ether);

        vm.warp(2);
        vm.expectRevert();
        router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            1 ether,
            1 ether,
            1 ether,
            address(this),
            1
        );
    }

    function testRemoveLiquidityEth() public {
        token0.approve(address(router), 10 ether);
        weth.approve(address(router), 10 ether);

        (uint256 amount0, uint256 amount1, uint256 liquidity) = router.addLiquidityETH{value: 1 ether}(
            address(token0),
            1 ether,
            1 ether,
            1 ether,
            address(this),
            block.timestamp + 1
        );

        uint256 prevAmount0 = token0.balanceOf(address(this));
        uint256 prevAmount1 = address(this).balance;

        address pair = factory.getPair(address(token0), address(weth));
        assertEq(IERC20(pair).balanceOf(address(this)), liquidity);
        IERC20(pair).approve(address(router), liquidity);

        router.removeLiquidityETH(
            address(token0),
            liquidity,
            0,
            0,
            address(this),
            block.timestamp + 1
        );
        UniswapV2Pair(pair).skim(address(this));

        assertEq(IERC20(pair).balanceOf(address(this)), 0);
        assertEq(token0.balanceOf(address(this)), prevAmount0 + amount0 - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
        assertEq(address(this).balance, prevAmount1 + amount1 - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());

        assertEq(token0.balanceOf(pair), UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
        assertEq(weth.balanceOf(address(pair)), UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
    }

    function testRemoveLiquidity() public {
        token0.approve(address(router), 1 ether);
        token1.approve(address(router), 1 ether);

        (uint256 amount0, uint256 amount1, uint256 liquidity) = router.addLiquidity(
            address(token0),
            address(token1),
            1 ether,
            1 ether,
            1 ether,
            1 ether,
            address(this),
            block.timestamp + 1
        );

        uint256 prevAmount0 = token0.balanceOf(address(this));
        uint256 prevAmount1 = token1.balanceOf(address(this));

        address pair = factory.getPair(address(token0), address(token1));
        assertEq(IERC20(pair).balanceOf(address(this)), liquidity);
        IERC20(pair).approve(address(router), liquidity);

        router.removeLiquidity(
            address(token0),
            address(token1),
            liquidity,
            0,
            0,
            address(this),
            block.timestamp + 1
        );
        UniswapV2Pair(pair).skim(address(this));

        assertEq(IERC20(pair).balanceOf(address(this)), 0);
        assertEq(token0.balanceOf(address(this)), prevAmount0 + amount0 - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
        assertEq(token1.balanceOf(address(this)), prevAmount1 + amount1 - UniswapV2Pair(pair).MINIMUM_LIQUIDITY());

        assertEq(token0.balanceOf(pair), UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
        assertEq(token1.balanceOf(pair), UniswapV2Pair(pair).MINIMUM_LIQUIDITY());
    }
}