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