perf: break early on proxy detection when realize can't get storage (ApeWorX#2514)

antazoey · web-flow · commit 7e5e462b1dcf · 2025-02-19T15:11:00.000Z
diff --git a/src/ape_ethereum/ecosystem.py b/src/ape_ethereum/ecosystem.py
@@ -505,8 +505,10 @@ def str_to_slot(text):
             try:
                 # TODO perf: use a batch call here when ape adds support
                 storage = self.provider.get_storage(address, slot)
-            except APINotImplementedError:
-                continue
+            except NotImplementedError:
+                # Break early on not-implemented error rather than attempting
+                # to try more proxy types.
+                break
 
             if sum(storage) == 0:
                 continue
diff --git a/tests/functional/geth/test_proxy.py b/tests/functional/geth/test_proxy.py
@@ -7,9 +7,6 @@
 
 @geth_process_test
 def test_standard_proxy(get_contract_type, owner, geth_contract, ethereum):
-    """
-    NOTE: Geth is used here because EthTester does not implement getting storage slots.
-    """
     _type = get_contract_type("eip1967")
     contract = ContractContainer(_type)
     target = geth_contract.address
@@ -22,9 +19,6 @@ def test_standard_proxy(get_contract_type, owner, geth_contract, ethereum):
 
 @geth_process_test
 def test_beacon_proxy(get_contract_type, geth_contract, owner, ethereum):
-    """
-    NOTE: Geth is used here because EthTester does not implement getting storage slots.
-    """
     _type = get_contract_type("beacon")
     beacon_contract = ContractContainer(_type)
     target = geth_contract.address
diff --git a/tests/functional/test_proxy.py b/tests/functional/test_proxy.py
@@ -1,4 +1,13 @@
+from typing import TYPE_CHECKING, Optional
+
+from eth_pydantic_types import HexBytes
+
+from ape.contracts.base import ContractContainer
 from ape_ethereum.proxies import ProxyType
+from ape_test.provider import LocalProvider
+
+if TYPE_CHECKING:
+    from ape.types import AddressType, BlockID
 
 """
 NOTE: Most proxy tests are in `geth/test_proxy.py`.
@@ -24,3 +33,49 @@ def test_minimal_proxy(ethereum, minimal_proxy_container, chain, owner):
     if isinstance(abi, list):
         assert abi == []
     # else: is messed up from other test (xdist).
+
+
+def test_provider_not_supports_get_storage(
+    get_contract_type, owner, vyper_contract_instance, ethereum, chain, networks
+):
+    """
+    The get storage slot RPC is required to detect this proxy, so it won't work
+    on EthTester provider. However, we can make sure that it doesn't try to
+    call `get_storage()` more than once.
+    """
+
+    class MyProvider(LocalProvider):
+        times_get_storage_was_called: int = 0
+
+        def get_storage(  # type: ignore[empty-body]
+            self, address: "AddressType", slot: int, block_id: Optional["BlockID"] = None
+        ) -> "HexBytes":
+            self.times_get_storage_was_called += 1
+            raise NotImplementedError()
+
+    my_provider = MyProvider(name="test", network=ethereum.local)
+    my_provider._web3 = chain.provider._web3
+
+    _type = get_contract_type("beacon")
+    beacon_contract = ContractContainer(_type)
+    target = vyper_contract_instance.address
+    beacon_instance = owner.deploy(beacon_contract, target)
+    beacon = beacon_instance.address
+
+    _type = get_contract_type("BeaconProxy")
+    contract = ContractContainer(_type)
+    contract_instance = owner.deploy(contract, beacon, HexBytes(""))
+
+    # Ensure not already cached.
+    if contract_instance.address in chain.contracts.proxy_infos:
+        del chain.contracts.proxy_infos[contract_instance.address]
+
+    init_provider = networks.active_provider
+    networks.active_provider = my_provider
+    try:
+        actual = ethereum.get_proxy_info(contract_instance.address)
+    finally:
+        networks.active_provider = init_provider
+
+    assert actual is None  # Because of provider.
+    assert my_provider.times_get_storage_was_called == 1
