fluxswap/test/FluxRouter.t.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.36;

import {Test} from "forge-std/Test.sol";
import {FluxFactory} from "../src/FluxFactory.sol";
import {FluxRouter} from "../src/FluxRouter.sol";
import {FluxPair} from "../src/FluxPair.sol";
import {TestErc20} from "../src/TestErc20.sol";

contract FluxRouterTest is Test {
    address private constant receiver = address(0xBEEF);

    FluxFactory private factory;
    FluxRouter private router;
    TestErc20 private tokenA;
    TestErc20 private tokenB;
    TestErc20 private tokenC;

    function setUp() public {
        vm.warp(1000);
        factory = new FluxFactory(address(this));
        router = new FluxRouter(address(factory));
        tokenA = new TestErc20("Alpha", "ALPHA", 1_000_000e18);
        tokenB = new TestErc20("Beta", "BETA", 1_000_000e18);
        tokenC = new TestErc20("Gamma", "GAMMA", 1_000_000e18);
        tokenA.approve(address(router), type(uint256).max);
        tokenB.approve(address(router), type(uint256).max);
        tokenC.approve(address(router), type(uint256).max);
    }

    function addAb(uint256 amountA, uint256 amountB) private returns (uint256 liquidity) {
        (,, liquidity) =
            router.addLiquidity(address(tokenA), address(tokenB), amountA, amountB, 0, 0, address(this), 2000);
    }

    function testAddLiquidityInitial() public {
        (uint256 amountA, uint256 amountB, uint256 liquidity) =
            router.addLiquidity(address(tokenA), address(tokenB), 100e18, 400e18, 0, 0, address(this), 2000);
        assertEq(amountA, 100e18);
        assertEq(amountB, 400e18);
        assertEq(liquidity, 199999999999999999000);

        address pair = factory.getPair(address(tokenA), address(tokenB));
        assertEq(FluxPair(pair).balanceOf(address(this)), 199999999999999999000);
        (uint112 reserveA, uint112 reserveB) = reservesOf(address(tokenA), address(tokenB));
        assertEq(reserveA, 100e18);
        assertEq(reserveB, 400e18);
    }

    function testAddLiquidityUsesOptimalAmount() public {
        addAb(100e18, 400e18);
        (uint256 amountA, uint256 amountB,) =
            router.addLiquidity(address(tokenA), address(tokenB), 50e18, 1000e18, 0, 0, address(this), 2000);
        assertEq(amountA, 50e18);
        assertEq(amountB, 200e18);
    }

    function testAddLiquidityRevertsOnBSlippage() public {
        addAb(100e18, 400e18);
        vm.expectRevert(bytes("FluxRouter: INSUFFICIENT_B_AMOUNT"));
        router.addLiquidity(address(tokenA), address(tokenB), 50e18, 1000e18, 0, 300e18, address(this), 2000);
    }

    function testRemoveLiquidityReturnsProportional() public {
        uint256 liquidity = addAb(100e18, 400e18);
        address pair = factory.getPair(address(tokenA), address(tokenB));
        FluxPair(pair).approve(address(router), liquidity);

        (uint256 amountA, uint256 amountB) =
            router.removeLiquidity(address(tokenA), address(tokenB), liquidity, 0, 0, receiver, 2000);

        assertEq(amountA, 99999999999999999500);
        assertEq(amountB, 399999999999999998000);
        assertEq(tokenA.balanceOf(receiver), 99999999999999999500);
        assertEq(tokenB.balanceOf(receiver), 399999999999999998000);
    }

    function testRemoveLiquidityRevertsOnSlippage() public {
        uint256 liquidity = addAb(100e18, 400e18);
        address pair = factory.getPair(address(tokenA), address(tokenB));
        FluxPair(pair).approve(address(router), liquidity);

        vm.expectRevert(bytes("FluxRouter: INSUFFICIENT_A_AMOUNT"));
        router.removeLiquidity(address(tokenA), address(tokenB), liquidity, 100e18, 0, receiver, 2000);
    }

    function testSwapExactTokensForTokensSingleHop() public {
        addAb(100e18, 400e18);
        uint256 expectedOut = 36264435755205965263;

        address[] memory path = new address[](2);
        path[0] = address(tokenA);
        path[1] = address(tokenB);

        uint256[] memory amounts = router.swapExactTokensForTokens(10e18, 0, path, receiver, 2000);

        assertEq(amounts[0], 10e18);
        assertEq(amounts[1], expectedOut);
        assertEq(tokenB.balanceOf(receiver), expectedOut);
        (uint112 reserveA, uint112 reserveB) = reservesOf(address(tokenA), address(tokenB));
        assertEq(reserveA, 110e18);
        assertEq(reserveB, 400e18 - expectedOut);
    }

    function testSwapTokensForExactTokensSingleHop() public {
        addAb(100e18, 400e18);
        uint256 expectedIn = 14328700386874910446;

        address[] memory path = new address[](2);
        path[0] = address(tokenA);
        path[1] = address(tokenB);

        uint256[] memory amounts = router.swapTokensForExactTokens(50e18, type(uint256).max, path, receiver, 2000);

        assertEq(amounts[0], expectedIn);
        assertEq(amounts[1], 50e18);
        assertEq(tokenB.balanceOf(receiver), 50e18);
    }

    function testMultiHopSwap() public {
        addAb(100e18, 200e18);
        router.addLiquidity(address(tokenB), address(tokenC), 300e18, 150e18, 0, 0, address(this), 2000);

        address[] memory path = new address[](3);
        path[0] = address(tokenA);
        path[1] = address(tokenB);
        path[2] = address(tokenC);

        uint256[] memory amounts = router.swapExactTokensForTokens(10e18, 0, path, receiver, 2000);

        assertEq(amounts[0], 10e18);
        assertEq(amounts[1], 18132217877602982631);
        assertEq(amounts[2], 8525187871197527614);
        assertEq(tokenC.balanceOf(receiver), 8525187871197527614);
    }

    function testSwapRevertsOnSlippage() public {
        addAb(100e18, 400e18);
        address[] memory path = new address[](2);
        path[0] = address(tokenA);
        path[1] = address(tokenB);

        vm.expectRevert(bytes("FluxRouter: INSUFFICIENT_OUTPUT_AMOUNT"));
        router.swapExactTokensForTokens(10e18, 40e18, path, receiver, 2000);
    }

    function testSwapRevertsOnExcessiveInput() public {
        addAb(100e18, 400e18);
        address[] memory path = new address[](2);
        path[0] = address(tokenA);
        path[1] = address(tokenB);

        vm.expectRevert(bytes("FluxRouter: EXCESSIVE_INPUT_AMOUNT"));
        router.swapTokensForExactTokens(50e18, 14e18, path, receiver, 2000);
    }

    function testSwapRevertsOnExpiredDeadline() public {
        addAb(100e18, 400e18);
        address[] memory path = new address[](2);
        path[0] = address(tokenA);
        path[1] = address(tokenB);

        vm.expectRevert(bytes("FluxRouter: EXPIRED"));
        router.swapExactTokensForTokens(10e18, 0, path, receiver, 999);
    }

    function reservesOf(address first, address second)
        private
        view
        returns (uint112 reserveFirst, uint112 reserveSecond)
    {
        address pair = factory.getPair(first, second);
        (uint112 reserve0, uint112 reserve1,) = FluxPair(pair).getReserves();
        (address token0,) = first < second ? (first, second) : (second, first);
        (reserveFirst, reserveSecond) = first == token0 ? (reserve0, reserve1) : (reserve1, reserve0);
    }
}