From fc5ccc9b9b1c64b5b99254e250e9553b539e6cfe Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Wed, 13 Oct 2021 00:38:05 -0400
Subject: [PATCH] Aave transfers

---
 mev_inspect/aave_liquidations.py | 32 +++++++++++++++++++++++++-------
 mev_inspect/schemas/transfers.py |  2 --
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/mev_inspect/aave_liquidations.py b/mev_inspect/aave_liquidations.py
index dc58673..db45a37 100644
--- a/mev_inspect/aave_liquidations.py
+++ b/mev_inspect/aave_liquidations.py
@@ -11,6 +11,10 @@ from mev_inspect.schemas.classified_traces import (
     Protocol,
 )
 
+<<<<<<< HEAD
+=======
+from mev_inspect.schemas.transfers import ERC20Transfer, aTokenTransfer, Transfer
+>>>>>>> Aave transfers
 from mev_inspect.schemas.liquidations import Liquidation
 from mev_inspect.transfers import get_transfer
 
@@ -84,15 +88,29 @@ def _get_payback_token_and_amount(
 ) -> Tuple[str, int]:
 
     """Look for and return liquidator payback from liquidation"""
+    child: ClassifiedTrace
+    child_transfer: Transfer
+
     for child in child_traces:
 
-        if child.classification == Classification.transfer:
+        if child.classification == Classification.transfer and isinstance(
+            child, DecodedCallTrace
+        ):
+            if liquidation.inputs["_receiveAToken"]:
 
-            child_transfer = ERC20Transfer.from_trace(child)
-            
-            if (
-                child_transfer.to_address == liquidator
-            ) and child.from_address in AAVE_CONTRACT_ADDRESSES:
-                return child_transfer.token_address, child_transfer.amount
+                child_transfer = aTokenTransfer.from_trace(child)
+
+                if (
+                    child_transfer.to_address == liquidator
+                ) and child.from_address in AAVE_CONTRACT_ADDRESSES:
+                    return child_transfer.token_address, child_transfer.amount
+
+            else:
+
+                child_transfer = ERC20Transfer.from_trace(child)
+
+                if (
+                    child_transfer.to_address == liquidator
+                ) and child.from_address in AAVE_CONTRACT_ADDRESSES:
 
     return liquidation.inputs["_collateral"], 0
diff --git a/mev_inspect/schemas/transfers.py b/mev_inspect/schemas/transfers.py
index 9719c2b..f39be0a 100644
--- a/mev_inspect/schemas/transfers.py
+++ b/mev_inspect/schemas/transfers.py
@@ -2,10 +2,8 @@ from typing import List
 
 from pydantic import BaseModel
 
-
 ETH_TOKEN_ADDRESS = "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE"
 
-
 class Transfer(BaseModel):
     block_number: int
     transaction_hash: str