tidy setup to perform the test directly over router

Pi-Squared-Inc · Sep 11, 2024 · 6e39cf0 · 6e39cf0
1 parent 95bd65f
commit 6e39cf0
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/test/regression/swaps.sol b/test/regression/swaps.sol
@@ -688,7 +688,8 @@ contract USDCMock {
 
 contract UniswapV2SwapTest {
 
-    UniswapV2Swap private _uni;
+    UniswapV2Swap private _uni;    
+    UniswapV2Router02 private _router;
     WETHMock private _weth;
     DAIMock private _dai;
     USDCMock private _usdc;
@@ -700,9 +701,9 @@ contract UniswapV2SwapTest {
         _uni = new UniswapV2Swap(address(_weth), address(_dai), address(_usdc));
         for (uint i = 0; i < 1000; i++) {
           testSwapSingleHopExactAmountIn();
-	}
+	    }
     }
-
+    
     function testSwapSingleHopExactAmountIn() public {
         uint256 wethAmount = 1e18;
         _weth.deposit{value: 2*wethAmount}();
@@ -725,19 +726,24 @@ contract UniswapV2SwapTest {
         _weth = new WETHMock();
         _dai = new DAIMock();
         _usdc = new USDCMock();
-        _uni = new UniswapV2Swap(address(_weth), address(_dai), address(_usdc));
+        _router = new UniswapV2Router02();
+
+        _router.set_local_pair(address(_weth), address(_dai));
+        _router.set_local_pair(address(_weth), address(_usdc));
+        _router.set_local_pair(address(_usdc), address(_dai));
+
         uint256 testAmount = 131072; // Hex: 0x20000
         _dai.mint(address(this), testAmount);
-        _dai.approve(address(_uni.router()), testAmount);
+        _dai.approve(address(_router), testAmount);
         _usdc.mint(address(this), testAmount);
-        _usdc.approve(address(_uni.router()), testAmount);
+        _usdc.approve(address(_router), testAmount);
 
-        _uni.router().addLiquidity(address(_dai), address(_usdc), 10000, 10000, 0, 0, address(this));
+        _router.addLiquidity(address(_dai), address(_usdc), 10000, 10000, 0, 0, address(this));
 
         assert(_dai.balanceOf(address(this)) == 121072);
         assert(_usdc.balanceOf(address(this)) == 121072);
-        assert(_dai.balanceOf(_uni.router().get_local_pair(address(_dai), address(_usdc))) == 10000);
-        assert(_usdc.balanceOf(_uni.router().get_local_pair(address(_dai), address(_usdc))) == 10000);
-        assert(UniswapV2Pair(_uni.router().get_local_pair(address(_dai), address(_usdc))).balanceOf(address(this)) == 9000);
+        assert(_dai.balanceOf(_router.get_local_pair(address(_dai), address(_usdc))) == 10000);
+        assert(_usdc.balanceOf(_router.get_local_pair(address(_dai), address(_usdc))) == 10000);
+        assert(UniswapV2Pair(_router.get_local_pair(address(_dai), address(_usdc))).balanceOf(address(this)) == 9000);
     }
 }