From bdcaaa9bf7b5b111c10df6e524774340f7c39ec2 Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Fri, 24 Sep 2021 17:44:58 -0400
Subject: [PATCH] Turned received amount logic to function

---
 mev_inspect/aave_liquidations.py | 31 +++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/mev_inspect/aave_liquidations.py b/mev_inspect/aave_liquidations.py
index 5f93601..d6efc17 100644
--- a/mev_inspect/aave_liquidations.py
+++ b/mev_inspect/aave_liquidations.py
@@ -1,4 +1,4 @@
-from typing import List, Dict
+from typing import List
 
 from mev_inspect.traces import get_child_traces
 from mev_inspect.schemas.classified_traces import (
@@ -26,7 +26,6 @@ def is_liquidator_payback(trace: ClassifiedTrace, liquidator: str) -> bool:
     """Finds liquidator payback """
 
     if isinstance(trace, DecodedCallTrace):
-
         if "recipient" in trace.inputs:
 
             if (
@@ -45,6 +44,20 @@ def is_liquidator_payback(trace: ClassifiedTrace, liquidator: str) -> bool:
     return False
 
 
+def get_received_amount(trace: DecodedCallTrace) -> int:
+
+    if "amount" in trace.inputs:
+        received_amount = int(trace.inputs["amount"])
+
+    elif "wad" in trace.inputs:
+        received_amount = int(trace.inputs["wad"])
+
+    else:
+        received_amount = 0
+
+    return received_amount
+
+
 def get_liquidations(
     traces: List[ClassifiedTrace],
 ) -> List[Liquidation]:
@@ -71,18 +84,8 @@ def get_liquidations(
 
                 if is_liquidator_payback(child, liquidator):
 
-                    assert isinstance(
-                        child.inputs, Dict
-                    )  # Necessary for mypy to indentify type
-
-                    if "amount" in child.inputs:
-                        received_amount = int(child.inputs["amount"])
-
-                    elif "wad" in child.inputs:
-                        received_amount = int(child.inputs["wad"])
-
-                    else:
-                        received_amount = 0
+                    assert isinstance(child, DecodedCallTrace)
+                    received_amount = get_received_amount(child)
 
             liquidations.append(
                 Liquidation(