From 2f1a9bc7514c9bcaf5730fa0aead273b7e100ecb Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Mon, 22 Nov 2021 12:35:23 -0500
Subject: [PATCH] Add helper for token_in_amount

---
 mev_inspect/classifiers/helpers.py       | 7 +++++++
 mev_inspect/classifiers/specs/zero_ex.py | 9 +++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/mev_inspect/classifiers/helpers.py b/mev_inspect/classifiers/helpers.py
index 58d6440..52b9219 100644
--- a/mev_inspect/classifiers/helpers.py
+++ b/mev_inspect/classifiers/helpers.py
@@ -84,3 +84,10 @@ def _filter_transfers(
         filtered_transfers.append(transfer)
 
     return filtered_transfers
+
+
+def get_amount_transferred_to_address(address: str, transfers: List[Transfer]) -> int:
+    for transfer in transfers:
+        if transfer.to_address == address:
+            return transfer.amount
+    return 0
diff --git a/mev_inspect/classifiers/specs/zero_ex.py b/mev_inspect/classifiers/specs/zero_ex.py
index 830c8b7..2b7d905 100644
--- a/mev_inspect/classifiers/specs/zero_ex.py
+++ b/mev_inspect/classifiers/specs/zero_ex.py
@@ -9,6 +9,7 @@ from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
     SwapClassifier,
 )
+from mev_inspect.classifiers.helpers import get_amount_transferred_to_address
 
 
 class ZeroExSwapClassifier(SwapClassifier):
@@ -40,11 +41,11 @@ class ZeroExSwapClassifier(SwapClassifier):
             raise NotImplementedError
 
         token_out_amount = trace.inputs["takerTokenFillAmount"]
-        contract_address = trace.to_address
+        token_in_amount = get_amount_transferred_to_address(
+            taker_address, child_transfers
+        )
 
-        for transfer in child_transfers:
-            if transfer.to_address == taker_address:
-                token_in_amount = transfer.amount
+        contract_address = trace.to_address
 
         return Swap(
             abi_name=trace.abi_name,