fix: custom err stuff

ape_foundry/provider.py

@@ -634,21 +634,10 @@ def _handle_execution_reverted(
 
         elif message.lower() == "execution reverted":
             message = TransactionError.DEFAULT_MESSAGE
-            if isinstance(exception, Web3ContractLogicError):
-                # Check for custom error.
-                trace = None
-                if "trace" in kwargs:
-                    trace = kwargs["trace"]
-
-                elif "txn" in kwargs:
-                    try:
-                        txn_hash = kwargs["txn"].txn_hash.hex()
-                        trace = self.get_transaction_trace(txn_hash)
-                    except Exception:
-                        pass
-
-                if trace is not None and (ret_value := trace.return_value):
-                    exception.message = ret_value
+            if isinstance(exception, Web3ContractLogicError) and (
+                msg := self._extract_custom_error(**kwargs)
+            ):
+                exception.message = msg
 
             return _handle_execution_reverted(exception, revert_message=message, **kwargs)
 
@@ -669,6 +658,26 @@ def _handle_execution_reverted(
 
         return VirtualMachineError(message, **kwargs)
 
+    # Abstracted for easier testing conditions.
+    def _extract_custom_error(self, **kwargs) -> str:
+        # Check for custom error.
+        trace = None
+        if "trace" in kwargs:
+            trace = kwargs["trace"]
+
+        elif "txn" in kwargs:
+            txn = kwargs["txn"]
+            try:
+                txn_hash = txn.txn_hash if isinstance(txn.txn_hash, str) else txn.txn_hash.hex()
+                trace = self.get_transaction_trace(txn_hash)
+            except Exception:
+                pass
+
+        if trace is not None and (revert_msg := trace.revert_message):
+            return revert_msg
+
+        return ""
+
     def set_block_gas_limit(self, gas_limit: int) -> bool:
         return self.make_request("evm_setBlockGasLimit", [hex(gas_limit)]) is True
 

tests/expected_traces.py

@@ -15,7 +15,7 @@
 │         333399998888882,
 │         234545457847457457458457457457
 │       \]
-│     \] \[\d+ gas\]
+│   \] \[\d+ gas\]
 ├── SYMBOL\.methodB1\(lolol="ice-cream", dynamo=345457847457457458457457457\) \[\d+ gas\]
 │   ├── ContractC\.getSomeList\(\) -> \[
 │   │     3425311345134513461345134534531452345,

tests/test_trace.py

@@ -4,6 +4,10 @@
 from typing import List
 
 import pytest
+from eth_utils import encode_hex
+from hexbytes import HexBytes
+
+from ape.exceptions import ContractLogicError
 from ape.utils import create_tempdir
 
 from .expected_traces import (
@@ -140,3 +144,40 @@ def assert_rich_output(rich_capture: List[str], expected: str):
     if expected_len > actual_len:
         rest = "\n".join(expected_lines[actual_len:])
         pytest.fail(f"Missing expected lines: {rest}")
+
+
+def test_extract_custom_error_trace_given(mocker, connected_provider):
+    trace = mocker.MagicMock()
+    trace.revert_message = "Unauthorized"
+    actual = connected_provider._extract_custom_error(trace=trace)
+    assert "Unauthorized" in actual
+
+
+def test_extract_custom_error_transaction_given(
+    connected_provider, vyper_contract_instance, not_owner
+):
+    with pytest.raises(ContractLogicError) as err:
+        vyper_contract_instance.setNumber(546, sender=not_owner, allow_fail=True)
+
+    actual = connected_provider._extract_custom_error(txn=err.value.txn)
+    assert actual == "!authorized"
+
+
+@pytest.mark.parametrize("tx_hash", ("0x0123", HexBytes("0x0123")))
+def test_extract_custom_error_transaction_given_trace_fails(connected_provider, mocker, tx_hash):
+    tx = mocker.MagicMock()
+    tx.txn_hash = tx_hash
+    tracker = []
+
+    def trace(txn_hash: str, *args, **kwargs):
+        tracker.append(txn_hash)
+        raise ValueError("Connection failed.")
+
+    connected_provider._get_transaction_trace = mocker.MagicMock()
+    connected_provider._get_transaction_trace.side_effect = trace
+
+    actual = connected_provider._extract_custom_error(txn=tx)
+    assert actual == ""
+
+    # Show failure was tracked
+    assert tracker[0] == HexBytes(tx.txn_hash).hex()