From e04027237f32c6180c41ac4e4e983ff1d6db36d2 Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Thu, 2 Sep 2021 15:53:04 -0400
Subject: [PATCH] removed optional from block and adjusted tests

---
 mev_inspect/classifiers/trace.py         |  4 +--
 mev_inspect/schemas/blocks.py            |  6 ++--
 mev_inspect/schemas/classified_traces.py |  5 ++-
 mev_inspect/schemas/swaps.py             | 24 ++++++++++++++
 scripts/poetry/inspect.py                |  3 +-
 tests/helpers.py                         | 40 ++++++++++++++++++------
 6 files changed, 64 insertions(+), 18 deletions(-)
 create mode 100644 mev_inspect/schemas/swaps.py

diff --git a/mev_inspect/classifiers/trace.py b/mev_inspect/classifiers/trace.py
index 2467346..00ddd65 100644
--- a/mev_inspect/classifiers/trace.py
+++ b/mev_inspect/classifiers/trace.py
@@ -7,7 +7,7 @@ from mev_inspect.schemas.classified_traces import (
     Classification,
     ClassifiedTrace,
     CallTrace,
-    ClassifiedCallTrace,
+    DecodedCallTrace,
     ClassifierSpec,
 )
 
@@ -71,7 +71,7 @@ class TraceClassifier:
                     signature, Classification.unknown
                 )
 
-                return ClassifiedCallTrace(
+                return DecodedCallTrace(
                     **trace.dict(),
                     trace_type=trace.type,
                     classification=classification,
diff --git a/mev_inspect/schemas/blocks.py b/mev_inspect/schemas/blocks.py
index 85ea3ce..33f9e46 100644
--- a/mev_inspect/schemas/blocks.py
+++ b/mev_inspect/schemas/blocks.py
@@ -46,9 +46,9 @@ class TraceType(Enum):
 
 class Trace(CamelModel):
     block_number: int
-    block_hash: Optional[str]
-    subtraces: Optional[int]
-    action: Optional[dict]
+    block_hash: str
+    action: dict
+    subtraces: int
     result: Optional[dict]
     trace_address: List[int]
     transaction_hash: Optional[str]
diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py
index 57a5500..3c1b094 100644
--- a/mev_inspect/schemas/classified_traces.py
+++ b/mev_inspect/schemas/classified_traces.py
@@ -51,6 +51,9 @@ class ClassifiedTrace(Trace):
     inputs: Optional[Dict[str, Any]]
     abi_name: Optional[str]
 
+    class Config:
+        validate_assignment = True
+
 
 class CallTrace(ClassifiedTrace):
 
@@ -58,7 +61,7 @@ class CallTrace(ClassifiedTrace):
     from_address: str
 
 
-class ClassifiedCallTrace(CallTrace):
+class DecodedCallTrace(CallTrace):
 
     inputs: Dict[str, Any]
     abi_name: str
diff --git a/mev_inspect/schemas/swaps.py b/mev_inspect/schemas/swaps.py
new file mode 100644
index 0000000..e87a6c8
--- /dev/null
+++ b/mev_inspect/schemas/swaps.py
@@ -0,0 +1,24 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from mev_inspect.schemas.classified_traces import Protocol
+
+
+class Swap(BaseModel):
+
+    abi_name: str
+    transaction_hash: str
+    block_number: int
+    trace_address: List[int]
+    pool_address: str
+    from_address: str
+    to_address: str
+    token_in_address: str
+    token_in_amount: int
+    token_out_address: str
+    token_out_amount: int
+
+    # Optional
+    protocol: Optional[Protocol]
+    error: Optional[str]
diff --git a/scripts/poetry/inspect.py b/scripts/poetry/inspect.py
index df6979b..b78a093 100644
--- a/scripts/poetry/inspect.py
+++ b/scripts/poetry/inspect.py
@@ -18,8 +18,7 @@ import click
 def inspect(block_number: str, rpc: str, cache: bool):
     check_call(
         [
-            "docker",
-            "compose",
+            "docker-compose",
             "exec",
             "mev-inspect",
             "python",
diff --git a/tests/helpers.py b/tests/helpers.py
index d27f57e..6f188f7 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -5,7 +5,7 @@ from mev_inspect.schemas.classified_traces import (
     Classification,
     ClassifiedTrace,
     CallTrace,
-    ClassifiedCallTrace,
+    DecodedCallTrace,
 )
 
 
@@ -17,6 +17,8 @@ def make_transfer_trace(
     to_address: str,
     token_address: str,
     amount: int,
+    action: dict = {},
+    subtraces: int = 0,
 ):
     return CallTrace(
         transaction_hash=transaction_hash,
@@ -30,6 +32,9 @@ def make_transfer_trace(
             "recipient": to_address,
             "amount": amount,
         },
+        block_hash=str(block_number),
+        action=action,
+        subtraces=subtraces,
     )
 
 
@@ -42,8 +47,12 @@ def make_swap_trace(
     abi_name: str,
     recipient_address: str,
     recipient_input_key: str,
+    action: dict = {},
+    subtraces: int = 0,
 ):
-    return ClassifiedCallTrace(
+    return DecodedCallTrace(
+        action=action,
+        subtraces=subtraces,
         transaction_hash=transaction_hash,
         block_number=block_number,
         type=TraceType.call,
@@ -53,27 +62,36 @@ def make_swap_trace(
         to_address=pool_address,
         inputs={recipient_input_key: recipient_address},
         abi_name=abi_name,
+        block_hash=str(block_number),
     )
 
 
 def make_unknown_trace(
-    block_number,
-    transaction_hash,
-    trace_address,
+    block_number: int,
+    transaction_hash: str,
+    trace_address: List[int],
+    action={},
+    subtraces=0,
 ):
+
     return ClassifiedTrace(
-        transaction_hash=transaction_hash,
         block_number=block_number,
-        type=TraceType.call,
+        transaction_hash=transaction_hash,
         trace_address=trace_address,
+        action=action,
+        subtraces=subtraces,
+        block_hash=str(block_number),
+        type=TraceType.call,
         classification=Classification.unknown,
     )
 
 
 def make_many_unknown_traces(
-    block_number,
-    transaction_hash,
-    trace_addresses,
+    block_number: int,
+    transaction_hash: str,
+    trace_addresses: List[List[int]],
+    action: dict = {},
+    subtraces: int = 0,
 ) -> List[ClassifiedTrace]:
 
     return [
@@ -81,6 +99,8 @@ def make_many_unknown_traces(
             block_number,
             transaction_hash,
             trace_address,
+            action,
+            subtraces,
         )
         for trace_address in trace_addresses
     ]