From b64d22c92b33fe0332b50b2448092210c5f65059 Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Wed, 18 Aug 2021 12:01:45 -0400
Subject: [PATCH] Enabled both stricter and previous object definition

---
 mev_inspect/schemas/classified_traces.py | 26 ++++++++++++++++++++++++
 tests/helpers.py                         |  5 ++---
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py
index 63cbea8..56a4c76 100644
--- a/mev_inspect/schemas/classified_traces.py
+++ b/mev_inspect/schemas/classified_traces.py
@@ -29,6 +29,32 @@ class ClassifiedTrace(BaseModel):
     trace_type: TraceType
     trace_address: List[int]
     classification: Classification
+    protocol: Optional[Protocol]
+    abi_name: Optional[str]
+    function_name: Optional[str]
+    function_signature: Optional[str]
+    inputs: Optional[Dict[str, Any]]
+    to_address: Optional[str]
+    from_address: Optional[str]
+    gas: Optional[int]
+    value: Optional[int]
+    gas_used: Optional[int]
+    error: Optional[str]
+
+    class Config:
+        json_encoders = {
+            # a little lazy but fine for now
+            # this is used for bytes value inputs
+            bytes: lambda b: b.hex(),
+        }
+
+
+class ClassifiedTraceStrict(BaseModel):
+    transaction_hash: str
+    block_number: int
+    trace_type: TraceType
+    trace_address: List[int]
+    classification: Classification
 
     class Config:
         json_encoders = {
diff --git a/tests/helpers.py b/tests/helpers.py
index 0f3677e..dcd15dd 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -4,7 +4,6 @@ from mev_inspect.schemas.blocks import TraceType
 from mev_inspect.schemas.classified_traces import (
     Classification,
     ClassifiedTrace,
-    ClassifiedCallData,
 )
 
 
@@ -17,7 +16,7 @@ def make_transfer_trace(
     token_address: str,
     amount: int,
 ):
-    return ClassifiedCallData(
+    return ClassifiedTrace(
         transaction_hash=transaction_hash,
         block_number=block_number,
         trace_type=TraceType.call,
@@ -42,7 +41,7 @@ def make_swap_trace(
     recipient_address: str,
     recipient_input_key: str,
 ):
-    return ClassifiedCallData(
+    return ClassifiedTrace(
         transaction_hash=transaction_hash,
         block_number=block_number,
         trace_type=TraceType.call,