From 46de116b5ce619ee55bd34fb904a2b33a96307d3 Mon Sep 17 00:00:00 2001
From: Mikko Ohtamaa <mikko@opensourcehacker.com>
Date: Fri, 12 Jan 2024 15:08:18 +0200
Subject: [PATCH] Adding guard tests

---
 eth_defi/simple_vault/__init__.py    |  6 ++
 eth_defi/simple_vault/transact.py    | 34 +++++++++++
 tests/guard/test_guard_uniswap_v2.py | 86 +++++++++++++++++++++-------
 3 files changed, 104 insertions(+), 22 deletions(-)
 create mode 100644 eth_defi/simple_vault/__init__.py
 create mode 100644 eth_defi/simple_vault/transact.py

diff --git a/eth_defi/simple_vault/__init__.py b/eth_defi/simple_vault/__init__.py
new file mode 100644
index 000000000..c195f7c90
--- /dev/null
+++ b/eth_defi/simple_vault/__init__.py
@@ -0,0 +1,6 @@
+"""Simple vault helpers.
+
+- Simple vault is a simplified vault implementation to test GuardV0 smart contract
+
+
+"""
\ No newline at end of file
diff --git a/eth_defi/simple_vault/transact.py b/eth_defi/simple_vault/transact.py
new file mode 100644
index 000000000..bdf1d120d
--- /dev/null
+++ b/eth_defi/simple_vault/transact.py
@@ -0,0 +1,34 @@
+from typing import Tuple
+
+from eth_typing import ChecksumAddress, HexStr
+from web3._utils.contracts import get_function_info, encode_abi
+from web3.contract.contract import ContractFunction
+
+
+def encode_simple_vault_transaction(func: ContractFunction) -> Tuple[ChecksumAddress, HexStr]:
+    """Encode a bound web3 function call as a simple vault transaction.
+
+    :param call:
+        Bound function prepared for a call.
+
+    :return:
+        Address, call data tuple.
+    """
+    assert isinstance(func, ContractFunction)
+
+    w3 = func.w3
+    contract_abi = func.contract_abi
+    fn_abi = func.abi
+    fn_identifier = func.function_identifier
+    args = func.args
+    fn_abi, fn_selector, fn_arguments = get_function_info(
+        # type ignored b/c fn_id here is always str b/c FallbackFn is handled above
+        fn_identifier,  # type: ignore
+        w3.codec,
+        contract_abi,
+        fn_abi,
+        args,
+    )
+    encoded = encode_abi(w3, fn_abi, fn_arguments, fn_selector)
+    return func.address, encoded
+
diff --git a/tests/guard/test_guard_uniswap_v2.py b/tests/guard/test_guard_uniswap_v2.py
index ab754d0a5..74a13bbee 100644
--- a/tests/guard/test_guard_uniswap_v2.py
+++ b/tests/guard/test_guard_uniswap_v2.py
@@ -5,8 +5,10 @@
 
 from eth_defi.abi import get_contract
 from eth_defi.deploy import deploy_contract
+from eth_defi.simple_vault.transact import encode_simple_vault_transaction
 from eth_defi.token import create_token
-from eth_defi.uniswap_v2.deployment import deploy_uniswap_v2_like, deploy_trading_pair
+from eth_defi.uniswap_v2.deployment import deploy_uniswap_v2_like, deploy_trading_pair, UniswapV2Deployment, FOREVER_DEADLINE
+from eth_defi.uniswap_v2.pair import fetch_pair_details, PairDetails
 
 
 @pytest.fixture
@@ -39,21 +41,14 @@ def deployer(web3) -> str:
 
 @pytest.fixture()
 def owner(web3) -> str:
-    """User account.
-
-    Do some account allocation for tests.
-    """
     return web3.eth.accounts[1]
 
 
 @pytest.fixture()
-def user_2(web3) -> str:
-    """User account.
-
-    Do some account allocation for tests.
-    """
+def asset_manager(web3) -> str:
     return web3.eth.accounts[2]
 
+
 @pytest.fixture()
 def usdc(web3, deployer) -> Contract:
     """Mock USDC token.
@@ -65,25 +60,72 @@ def usdc(web3, deployer) -> Contract:
 
 
 @pytest.fixture()
-def uniswap_v2(web3: Web3, usdc: Contract, deployer: str):
+def uniswap_v2(web3: Web3, usdc: Contract, deployer: str) -> UniswapV2Deployment:
     """Deploy mock Uniswap v2."""
     uniswap_v2 = deploy_uniswap_v2_like(web3, deployer)
-    pair_address = deploy_trading_pair(
-        web3,
-        deployer,
-        uniswap_v2,
-        uniswap_v2.weth,
-        usdc,
-        0,  # 10 ETH liquidity
-        0,  # 17000 USDC liquidity
-    )
 
 
 @pytest.fixture()
-def vault(web3: Web3, usdc: Contract, deployer: str) -> Contract:
+def vault(web3: Web3, usdc: Contract, deployer: str, owner: str, asset_manager: str) -> Contract:
     """Deploy mock Uniswap v2."""
     uniswap_v2 = deploy_uniswap_v2_like(web3, deployer)
+    weth = uniswap_v2.weth
     vault = deploy_contract(web3, "guard/SimpleVaultV0.json", deployer)
+    vault.functions.transferOwnership(owner).transact({"from": deployer})
+    vault.functions.updateAssetManager(asset_manager).transact({"from": owner})
     guard = get_contract(web3, "guard/GuardV0.json", vault.functions.guard.call())
-    guard.functions.
+    guard.functions.whitelistUniswapV2Router(uniswap_v2.router.address).transact({"from": deployer})
+    guard.functions.whitelistToken(usdc.address)
+    guard.functions.whitelistToken(weth.address)
     return vault
+
+
+@pytest.fixture()
+def weth(uniswap_v2) -> Contract:
+    return uniswap_v2.weth
+
+
+@pytest.fixture()
+def weth_usdc_pair(uniswap_v2, weth, usdc, deployer) -> PairDetails:
+    pair_address = deploy_trading_pair(
+        web3,
+        deployer,
+        uniswap_v2,
+        weth,
+        usdc,
+        10 * 10**18,  # 10 ETH liquidity
+        17_000 * 10**6,  # 17000 USDC liquidity
+    )
+    return fetch_pair_details(web3, pair_address)
+
+
+def test_guard_can_trade_uniswap_v2(
+    uniswap_v2: UniswapV2Deployment,
+    owner: str,
+    asset_manager: str,
+    weth: Contract,
+    usdc: Contract,
+    vault: Contract
+):
+    usdc_amount = 10_000 ** 10**6
+    usdc.functions.transfer(vault.address, 10_000).transact({"from": deployer})
+
+    path = [usdc.address, weth.address]
+
+    trade_call = uniswap_v2.router.functions.swapExactTokensForTokens(
+        usdc_amount,
+        0,
+        path,
+        vault.address,
+        FOREVER_DEADLINE,
+    )
+
+    target, call_data = encode_simple_vault_transaction(trade_call)
+    vault.functions.performCall(target, call_data).transact({"from": asset_manager})
+
+
+
+
+
+
+