feat: trading with V3 now working!

ErikBjare · ErikBjare · commit a4091f367df9 · 2021-05-17T14:43:53.000+02:00
diff --git a/tests/test_uniswap.py b/tests/test_uniswap.py
@@ -37,17 +37,13 @@ def test_assets(client: Uniswap):
     """
     tokens = client._get_token_addresses()
 
-    for token_name, amount in [("DAI", 10 ** 18), ("USDC", 10 ** 6)]:
+    for token_name, amount in [("DAI", 100 * 10 ** 18), ("USDC", 100 * 10 ** 6)]:
         token_addr = tokens[token_name]
-        price = client.get_eth_token_output_price(_str_to_addr(token_addr), amount)
+        price = client.get_eth_token_output_price(token_addr, amount)
         logger.info(f"Cost of {amount} {token_name}: {price}")
         logger.info("Buying...")
 
-        client.approve(tokens["WETH"])
-        # client.approve(tokens[token_name])
-
-        # FIXME: Maybe trading ETH like this isn't supported on V3?
-        tx = client.make_trade_output(tokens["WETH"], tokens[token_name], amount)
+        tx = client.make_trade_output(tokens["ETH"], token_addr, amount)
         client.w3.eth.waitForTransactionReceipt(tx)
 
 
@@ -271,11 +267,11 @@ def test_remove_liquidity(
         "input_token, output_token, qty, recipient, expectation",
         [
             # ETH -> Token
-            (eth, bat, 1_000_000_000 * ONE_WEI, None, does_not_raise),
+            (eth, dai, 1_000_000_000 * ONE_WEI, None, does_not_raise),
             # Token -> Token
-            # (bat, dai, 1_000_000_000 * ONE_WEI, None, does_not_raise),
+            (dai, usdc, 1_000_000_000 * ONE_WEI, None, does_not_raise),
             # Token -> ETH
-            (bat, eth, 1_000_000 * ONE_WEI, None, does_not_raise),
+            (usdc, eth, 1_000_000 * ONE_WEI, None, does_not_raise),
             # (eth, bat, 0.00001 * ONE_ETH, ZERO_ADDRESS, does_not_raise),
             # (bat, eth, 0.00001 * ONE_ETH, ZERO_ADDRESS, does_not_raise),
             # (dai, bat, 0.00001 * ONE_ETH, ZERO_ADDRESS, does_not_raise),
@@ -309,10 +305,9 @@ def test_make_trade(
         "input_token, output_token, qty, recipient, expectation",
         [
             # ETH -> Token
-            (eth, bat, 1_000_000 * ONE_WEI, None, does_not_raise),
+            (eth, dai, 1_000_000 * ONE_WEI, None, does_not_raise),
             # Token -> Token
-            # No liquidity in V3
-            # (bat, dai, 1_000_000 * ONE_WEI, None, does_not_raise),
+            (dai, usdc, 1_000_000 * ONE_WEI, None, does_not_raise),
             # Token -> ETH
             (dai, eth, 1_000_000 * ONE_WEI, None, does_not_raise),
             # FIXME: These should probably be uncommented eventually
diff --git a/uniswap/uniswap.py b/uniswap/uniswap.py
@@ -59,15 +59,13 @@ def check_approval(method: Callable) -> Callable:
 
     @functools.wraps(method)
     def approved(self: Any, *args: Any, **kwargs: Any) -> Any:
-        eth_addr = ETH_ADDRESS if self.version < 3 else WETH9_ADDRESS
-
         # Check to see if the first token is actually ETH
-        token = args[0] if args[0] != eth_addr else None
+        token = args[0] if args[0] != ETH_ADDRESS else None
         token_two = None
 
         # Check second token, if needed
         if method.__name__ == "make_trade" or method.__name__ == "make_trade_output":
-            token_two = args[1] if args[1] in [eth_addr] else None
+            token_two = args[1] if args[1] != ETH_ADDRESS else None
 
         # Approve both tokens, if needed
         if token:
@@ -654,7 +652,8 @@ def make_trade_output(
             )
         elif output_token == ETH_ADDRESS:
             qty = Wei(qty)
-            return self._token_to_eth_swap_output(input_token, qty, recipient)
+            # No idea why mypy complains here...
+            return self._token_to_eth_swap_output(input_token, qty, recipient)  # type: ignore
         else:
             return self._token_to_token_swap_output(
                 input_token, output_token, qty, recipient
@@ -678,7 +677,8 @@ def _eth_to_token_swap_input(
                 func_params.append(recipient)
                 function = token_funcs.ethToTokenTransferInput(*func_params)
             return self._build_and_send_tx(function, tx_params)
-        else:
+
+        elif self.version == 2:
             if recipient is None:
                 recipient = self.address
             amount_out_min = int(
@@ -694,6 +694,12 @@ def _eth_to_token_swap_input(
                 ),
                 self._get_tx_params(qty),
             )
+        elif self.version == 3:
+            return self._token_to_token_swap_input(
+                self.get_weth_address(), output_token, qty, recipient
+            )
+        else:
+            raise ValueError
 
     def _token_to_eth_swap_input(
         self, input_token: AddressLike, qty: int, recipient: Optional[AddressLike]
@@ -713,7 +719,7 @@ def _token_to_eth_swap_input(
                 func_params.append(recipient)
                 function = token_funcs.tokenToEthTransferInput(*func_params)
             return self._build_and_send_tx(function)
-        else:
+        elif self.version == 2:
             if recipient is None:
                 recipient = self.address
             amount_out_min = int(
@@ -729,6 +735,12 @@ def _token_to_eth_swap_input(
                     self._deadline(),
                 ),
             )
+        elif self.version == 3:
+            return self._token_to_token_swap_input(
+                input_token, self.get_weth_address(), qty, recipient
+            )
+        else:
+            raise ValueError
 
     def _token_to_token_swap_input(
         self,
