From d33179d2671e1546e9a22365beec69cfd87145d1 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Fri, 30 Jul 2021 14:13:24 -0400
Subject: [PATCH] Change NestedTrace to use ClassifiedTraces

---
 mev_inspect/schemas/__init__.py          |  2 +-
 mev_inspect/schemas/blocks.py            |  6 ------
 mev_inspect/schemas/classified_traces.py |  8 ++++++++
 mev_inspect/traces.py                    | 10 +++++-----
 tests/trace_test.py                      | 18 ++++++------------
 5 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/mev_inspect/schemas/__init__.py b/mev_inspect/schemas/__init__.py
index b4319d2..442724e 100644
--- a/mev_inspect/schemas/__init__.py
+++ b/mev_inspect/schemas/__init__.py
@@ -1,2 +1,2 @@
 from .abi import ABI
-from .blocks import Block, NestedTrace, Trace, TraceType
+from .blocks import Block, Trace, TraceType
diff --git a/mev_inspect/schemas/blocks.py b/mev_inspect/schemas/blocks.py
index b6e3635..d154029 100644
--- a/mev_inspect/schemas/blocks.py
+++ b/mev_inspect/schemas/blocks.py
@@ -68,9 +68,3 @@ class Block(Web3Model):
         return [trace for trace in self.traces if trace.transaction_hash == hash]
 
 
-class NestedTrace(BaseModel):
-    trace: Trace
-    subtraces: List["NestedTrace"]
-
-
-NestedTrace.update_forward_refs()
diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py
index 3394550..645413b 100644
--- a/mev_inspect/schemas/classified_traces.py
+++ b/mev_inspect/schemas/classified_traces.py
@@ -50,3 +50,11 @@ class ClassifierSpec(BaseModel):
     protocol: Optional[Protocol] = None
     valid_contract_addresses: Optional[List[str]] = None
     classifications: Dict[str, Classification] = {}
+
+
+class NestedTrace(BaseModel):
+    trace: ClassifiedTrace
+    subtraces: List["NestedTrace"]
+
+
+NestedTrace.update_forward_refs()
diff --git a/mev_inspect/traces.py b/mev_inspect/traces.py
index 792897e..052d8be 100644
--- a/mev_inspect/traces.py
+++ b/mev_inspect/traces.py
@@ -1,10 +1,10 @@
 from itertools import groupby
 from typing import Iterable, List
 
-from mev_inspect.schemas import Trace, NestedTrace
+from mev_inspect.schemas.classified_traces import ClassifiedTrace, NestedTrace
 
 
-def as_nested_traces(traces: Iterable[Trace]) -> List[NestedTrace]:
+def as_nested_traces(traces: Iterable[ClassifiedTrace]) -> List[NestedTrace]:
     nested_traces = []
 
     sorted_by_transaction_hash = sorted(traces, key=_get_transaction_hash)
@@ -20,7 +20,7 @@ def _get_transaction_hash(trace) -> str:
     return trace.transaction_hash
 
 
-def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrace]:
+def _as_nested_traces_by_transaction(traces: Iterable[ClassifiedTrace]) -> List[NestedTrace]:
     """
     Turns a list of Traces into a a tree of NestedTraces
     using their trace addresses
@@ -35,7 +35,7 @@ def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrac
     nested_traces = []
 
     parent = None
-    children: List[Trace] = []
+    children: List[ClassifiedTrace] = []
 
     sorted_traces = sorted(traces, key=lambda t: t.trace_address)
 
@@ -70,7 +70,7 @@ def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrac
     return nested_traces
 
 
-def _is_subtrace(trace: Trace, parent: Trace):
+def _is_subtrace(trace: ClassifiedTrace, parent: ClassifiedTrace):
     parent_trace_length = len(parent.trace_address)
 
     if len(trace.trace_address) > parent_trace_length:
diff --git a/tests/trace_test.py b/tests/trace_test.py
index dd415aa..d5478c5 100644
--- a/tests/trace_test.py
+++ b/tests/trace_test.py
@@ -1,7 +1,8 @@
 import unittest
 from typing import List
 
-from mev_inspect.schemas import Trace, TraceType, NestedTrace
+from mev_inspect.schemas import TraceType
+from mev_inspect.schemas.classified_traces import ClassifiedTrace, Classification, NestedTrace
 from mev_inspect.traces import as_nested_traces
 
 
@@ -82,20 +83,13 @@ class TestTraces(unittest.TestCase):
 def build_trace_at_address(
     transaction_hash: str,
     trace_address: List[int],
-) -> Trace:
-    return Trace(
-        # real values
+) -> ClassifiedTrace:
+    return ClassifiedTrace(
         transaction_hash=transaction_hash,
         trace_address=trace_address,
-        # placeholders
-        action={},
-        block_hash="",
         block_number=DEFAULT_BLOCK_NUMBER,
-        result=None,
-        subtraces=0,
-        transaction_position=None,
-        type=TraceType.call,
-        error=None,
+        trace_type=TraceType.call,
+        classification=Classification.unknown,
     )