Adjust AAVE and include prices fromm #231

mev_inspect/classifiers/specs/aave.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional
 
 from mev_inspect.schemas.classifiers import (
     ClassifiedTrace,
@@ -21,23 +21,28 @@ class AaveLiquidationClassifier(LiquidationClassifier):
     ) -> Optional[Liquidation]:
 
         liquidator = liquidation_trace.from_address
+        liquidated = liquidation_trace.inputs["_user"]
 
-        (debt_token_address, debt_purchase_amount) = _get_debt_data(
-            liquidation_trace, child_transfers, liquidator
-        )
+        debt_token_address = liquidation_trace.inputs["_reserve"]
+        received_token_address = liquidation_trace.inputs["_collateral"]
 
-        if debt_purchase_amount == 0:
-            return None
+        debt_purchase_amount = None
+        received_amount = None
 
-        (received_token_address, received_amount) = _get_received_data(
-            liquidation_trace, child_transfers, liquidator
-        )
+        debt_transfer = _get_debt_transfer(liquidator, child_transfers)
 
-        if received_amount == 0:
-            return None
+        received_transfer = _get_received_transfer(liquidator, child_transfers)
+
+        if debt_transfer is not None and received_transfer is not None:
+
+            debt_token_address = debt_transfer.token_address
+            debt_purchase_amount = debt_transfer.amount
+
+            received_token_address = received_transfer.token_address
+            received_amount = received_transfer.amount
 
             return Liquidation(
-            liquidated_user=liquidation_trace.inputs["_user"],
+                liquidated_user=liquidated,
                 debt_token_address=debt_token_address,
                 liquidator_user=liquidator,
                 debt_purchase_amount=debt_purchase_amount,
@@ -50,6 +55,9 @@ class AaveLiquidationClassifier(LiquidationClassifier):
                 error=liquidation_trace.error,
             )
 
+        else:
+            return None
+
 
 class AaveTransferClassifier(TransferClassifier):
     @staticmethod
@@ -84,30 +92,25 @@ ATOKENS_SPEC = ClassifierSpec(
 AAVE_CLASSIFIER_SPECS: List[ClassifierSpec] = [AAVE_SPEC, ATOKENS_SPEC]
 
 
-def _get_received_data(
-    liquidation_trace: DecodedCallTrace,
-    child_transfers: List[Transfer],
-    liquidator: str,
-) -> Tuple[str, int]:
+def _get_received_transfer(
+    liquidator: str, child_transfers: List[Transfer]
+) -> Optional[Transfer]:
+    """Get transfer from AAVE to liquidator"""
 
-    """Look for and return liquidator payback from liquidation"""
     for transfer in child_transfers:
-
         if transfer.to_address == liquidator:
-            return transfer.token_address, transfer.amount
+            return transfer
 
-    return liquidation_trace.inputs["_collateral"], 0
+    return None
 
 
-def _get_debt_data(
-    liquidation_trace: DecodedCallTrace,
-    child_transfers: List[Transfer],
-    liquidator: str,
-) -> Tuple[str, int]:
+def _get_debt_transfer(
+    liquidator: str, child_transfers: List[Transfer]
+) -> Optional[Transfer]:
     """Get transfer from liquidator to AAVE"""
 
     for transfer in child_transfers:
         if transfer.from_address == liquidator:
-            return transfer.token_address, transfer.amount
+            return transfer
 
-    return liquidation_trace.inputs["_reserve"], 0
+    return None

mev_inspect/classifiers/specs/compound.py

@@ -9,6 +9,7 @@ from mev_inspect.schemas.classifiers import (
     SeizeClassifier,
 )
 from mev_inspect.schemas.liquidations import Liquidation
+from mev_inspect.schemas.prices import CETH_TOKEN_ADDRESS
 from mev_inspect.schemas.traces import Protocol
 from mev_inspect.schemas.transfers import Transfer
 
@@ -23,11 +24,19 @@ class CompoundLiquidationClassifier(LiquidationClassifier):
 
         liquidator = liquidation_trace.from_address
         liquidated = liquidation_trace.inputs["borrower"]
+
         debt_token_address = liquidation_trace.to_address
-        debt_purchase_amount = liquidation_trace.inputs["repayAmount"]
         received_token_address = liquidation_trace.inputs["cTokenCollateral"]
+
+        debt_purchase_amount = None
         received_amount = None
 
+        if debt_token_address == CETH_TOKEN_ADDRESS:
+            debt_purchase_amount = liquidation_trace.value
+
+        else:
+            debt_purchase_amount = liquidation_trace.inputs["repayAmount"]
+
         debt_transfer = _get_debt_transfer(liquidator, child_transfers)
 
         received_transfer = _get_received_transfer(liquidator, child_transfers)

mev_inspect/schemas/prices.py

@@ -53,6 +53,7 @@ COINGECKO_ID_BY_ADDRESS = {
 }
 
 
+
 class Price(BaseModel):
     token_address: str
     usd_price: float