fix: address PR feedback

ApeWorX · Feb 18, 2025 · 9eeaeea · 9eeaeea
1 parent e17fa60
commit 9eeaeea
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 8 deletions.
diff --git a/src/ape/managers/_contractscache.py b/src/ape/managers/_contractscache.py
@@ -577,19 +577,22 @@ def get(
                     if fetch_from_explorer
                     else None
                 )
-                contract_type_to_cache = None
-
-                if proxy_contract_type:
-                    contract_type_to_cache = _get_combined_contract_type(
+                if proxy_contract_type is not None and implementation_contract_type is not None:
+                    combined_contract = _get_combined_contract_type(
                         proxy_contract_type, proxy_info, implementation_contract_type
                     )
+                    self.contract_types[address_key] = combined_contract
+                    return combined_contract
+
                 elif implementation_contract_type is not None:
                     contract_type_to_cache = implementation_contract_type
-
-                if contract_type_to_cache is not None:
-                    self.contract_types[address_key] = contract_type_to_cache
+                    self.contract_types[address_key] = implementation_contract_type
                     return contract_type_to_cache
 
+                elif proxy_contract_type is not None:
+                    self.contract_types[address_key] = proxy_contract_type
+                    return proxy_contract_type
+
             # Also gets cached to disk for faster lookup next time.
             if fetch_from_explorer:
                 contract_type = self._get_contract_type_from_explorer(address_key)

diff --git a/tests/functional/test_contracts_cache.py b/tests/functional/test_contracts_cache.py
@@ -474,7 +474,7 @@ def test_cache_non_checksum_address(chain, vyper_contract_instance):
     assert chain.contracts[vyper_contract_instance.address] == vyper_contract_instance.contract_type
 
 
-def test_get_when_proxy(chain, owner, minimal_proxy_container):
+def test_get_when_proxy(chain, owner, minimal_proxy_container, vyper_contract_instance):
     placeholder = "0xBEbeBeBEbeBebeBeBEBEbebEBeBeBebeBeBebebe"
     if placeholder in chain.contracts:
         del chain.contracts[placeholder]
@@ -486,6 +486,27 @@ def test_get_when_proxy(chain, owner, minimal_proxy_container):
     assert actual == minimal_proxy.contract_type
 
 
+def test_get_when_proxy_but_implementation_missing(chain, owner, vyper_contract_container):
+    """
+    Proxy is cached but implementation is missing.
+    """
+    placeholder = vyper_contract_container.deploy(1001, sender=owner)
+    assert chain.contracts[placeholder.address]  # This must be cached!
+
+    proxy_container = _make_minimal_proxy(placeholder.address)
+    minimal_proxy = owner.deploy(proxy_container, sender=owner)
+    chain.provider.network.__dict__["explorer"] = None  # Ensure no explorer, messes up test.
+
+    if minimal_proxy.address in chain.contracts:
+        # Delete the proxy but make sure it does not delete the implementation!
+        # (which it normally does here).
+        del chain.contracts[minimal_proxy.address]
+        chain.contracts[placeholder.address] = placeholder
+
+    actual = chain.contracts.get(minimal_proxy.address)
+    assert actual == minimal_proxy.contract_type
+
+
 def test_get_pass_along_proxy_info(chain, owner, minimal_proxy_container, ethereum):
     placeholder = "0xBEbeBeBEbeBebeBeBEBEbebEBeBeBebeBeBebebe"
     if placeholder in chain.contracts: