From 58c43312b3bc5f78bc569652e45d0834c1c65006 Mon Sep 17 00:00:00 2001
From: antazoey <jules@apeworx.io>
Date: Mon, 10 Jun 2024 08:03:58 -0500
Subject: [PATCH] fix: handle enum types more in trace enrichment (#2121)

---
 src/ape_ethereum/ecosystem.py      |  6 +++++
 tests/functional/test_ecosystem.py | 38 ++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/src/ape_ethereum/ecosystem.py b/src/ape_ethereum/ecosystem.py
index faaa6541a9..1b57bdccba 100644
--- a/src/ape_ethereum/ecosystem.py
+++ b/src/ape_ethereum/ecosystem.py
@@ -1016,6 +1016,7 @@ def enrich_trace(self, trace: TraceAPI, **kwargs) -> TraceAPI:
             kwargs["sender"] = sender
 
         # Get the un-enriched calltree.
+        # NOTE: Using JSON mode so Enums are all str types.
         data = trace.get_calltree().model_dump(mode="json", by_alias=True)
 
         if isinstance(trace, TransactionTrace):
@@ -1044,7 +1045,12 @@ def _enrich_calltree(self, call: dict, **kwargs) -> dict:
         kwargs["use_symbol_for_tokens"] = kwargs.get(
             "use_symbol_for_tokens", default_symbol_for_tokens
         )
+
+        # Handle if for some reason this is still an Enum.
         call_type = call.get("call_type", "")
+        if call_type and not isinstance(call_type, str):
+            call["call_type"] = call_type = call_type.value
+
         is_create = "CREATE" in call_type
 
         # Enrich sub-calls first.
diff --git a/tests/functional/test_ecosystem.py b/tests/functional/test_ecosystem.py
index 998804a478..f91abd70b4 100644
--- a/tests/functional/test_ecosystem.py
+++ b/tests/functional/test_ecosystem.py
@@ -1,4 +1,5 @@
 import copy
+import re
 from typing import Any, ClassVar, cast
 
 import pytest
@@ -6,11 +7,13 @@
 from eth_typing import HexAddress, HexStr
 from ethpm_types import ContractType, ErrorABI
 from ethpm_types.abi import ABIType, EventABI, MethodABI
+from evm_trace import CallTreeNode, CallType
 
 from ape.api.networks import LOCAL_NETWORK_NAME, NetworkAPI
 from ape.exceptions import CustomError, DecodingError, NetworkError, NetworkNotFoundError
 from ape.types import AddressType
 from ape.utils import DEFAULT_LOCAL_TRANSACTION_ACCEPTANCE_TIMEOUT
+from ape_ethereum import TransactionTrace
 from ape_ethereum.ecosystem import BLUEPRINT_HEADER, BaseEthereumConfig, Block
 from ape_ethereum.transactions import (
     DynamicFeeTransaction,
@@ -1033,3 +1036,38 @@ def test_decode_custom_error_selector_not_found(chain, ethereum):
     tx = ethereum.create_transaction()
     actual = ethereum.decode_custom_error(data, addr, txn=tx)
     assert actual is None
+
+
+def test_enrich_trace(ethereum, vyper_contract_instance, owner):
+    tx = vyper_contract_instance.setNumber(96247783, sender=owner)
+    trace = TransactionTrace(transaction_hash=tx.txn_hash)
+    actual = ethereum.enrich_trace(trace)
+    assert isinstance(actual, TransactionTrace)
+    assert actual._enriched_calltree is not None
+    assert re.match(r"VyperContract\.setNumber\(num=\d*\) \[\d* gas]", repr(actual))
+
+
+def test_enrich_trace_handles_call_type_enum(ethereum, vyper_contract_instance, owner):
+    """
+    Testing a custom trace who's call tree uses an Enum type instead of a str.
+    """
+
+    class PluginTxTrace(TransactionTrace):
+        def get_calltree(self) -> CallTreeNode:
+            call = super().get_calltree()
+            # Force enum value instead of str.
+            call.call_type = CallType.CALL
+            return call
+
+    tx = vyper_contract_instance.setNumber(96247783, sender=owner)
+    trace = PluginTxTrace(transaction_hash=tx.txn_hash)
+    actual = ethereum.enrich_trace(trace)
+    assert isinstance(actual, PluginTxTrace)
+    assert actual._enriched_calltree is not None
+    assert re.match(r"VyperContract\.setNumber\(num=\d*\) \[\d* gas]", repr(actual))
+
+    # Hook into helper method hackily with an enum.
+    # Notice the mode is Python and not JSON here.
+    call = trace.get_calltree().model_dump(by_alias=True)
+    actual = ethereum._enrich_calltree(call)
+    assert actual["call_type"] == CallType.CALL.value