From ea8325d67b89b72a6d2ca0548606215f6d26421d Mon Sep 17 00:00:00 2001
From: antazoey <antyzoa@gmail.com>
Date: Tue, 28 Jan 2025 10:40:27 -0600
Subject: [PATCH] fix: dont cache when block id givne

---
 src/ape/managers/chain.py           | 12 ++++++++----
 tests/functional/conftest.py        |  2 +-
 tests/functional/geth/test_chain.py | 18 ++++++++++++++++++
 3 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 tests/functional/geth/test_chain.py

diff --git a/src/ape/managers/chain.py b/src/ape/managers/chain.py
index b5296a698d..c4e640554e 100644
--- a/src/ape/managers/chain.py
+++ b/src/ape/managers/chain.py
@@ -970,9 +970,13 @@ def get_code(
         self, address: AddressType, block_id: Optional["BlockID"] = None
     ) -> "ContractCode":
         network = self.provider.network
-        if network.is_dev:
-            # Avoid caching when dev, as you can manipulate the chain more
-            # (and there is isolation).
+
+        # Two reasons to avoid caching:
+        #   1. dev networks - chain isolation makes this mess up
+        #   2. specifying block_id= kwarg - likely checking if code
+        #      exists at the time and shouldn't use cache.
+        skip_cache = network.is_dev or block_id is not None
+        if skip_cache:
             return self.provider.get_code(address, block_id=block_id)
 
         self._code.setdefault(network.ecosystem.name, {})
@@ -980,7 +984,7 @@ def get_code(
         if address in self._code[network.ecosystem.name][network.name]:
             return self._code[network.ecosystem.name][network.name][address]
 
-        # Get from RPC for the first time.
+        # Get from RPC for the first time AND use cache.
         code = self.provider.get_code(address)
         self._code[network.ecosystem.name][network.name][address] = code
         return code
diff --git a/tests/functional/conftest.py b/tests/functional/conftest.py
index 85c013a710..20f1557459 100644
--- a/tests/functional/conftest.py
+++ b/tests/functional/conftest.py
@@ -691,7 +691,7 @@ def create_mock_sepolia(ethereum, eth_tester_provider, vyper_contract_instance):
     @contextmanager
     def fn():
         # Ensuring contract exists before hack.
-        # This allow the network to be past genesis which is more realistic.
+        # This allows the network to be past genesis which is more realistic.
         _ = vyper_contract_instance
         eth_tester_provider.network.name = "sepolia"
         yield eth_tester_provider.network
diff --git a/tests/functional/geth/test_chain.py b/tests/functional/geth/test_chain.py
new file mode 100644
index 0000000000..09bd82ab41
--- /dev/null
+++ b/tests/functional/geth/test_chain.py
@@ -0,0 +1,18 @@
+from tests.conftest import geth_process_test
+
+
+@geth_process_test
+def test_get_code(mocker, chain, geth_contract, mock_sepolia):
+    # NOTE: Using mock_sepolia because code doesn't get cached in local networks.
+    actual = chain.get_code(geth_contract.address)
+    expected = chain.provider.get_code(geth_contract.address)
+    assert actual == expected
+
+    # Ensure uses cache (via not using provider).
+    provider_spy = mocker.spy(chain.provider.web3.eth, "get_code")
+    _ = chain.get_code(geth_contract.address)
+    assert provider_spy.call_count == 0
+
+    # block_id test, cache should interfere.
+    actual_2 = chain.get_code(geth_contract.address, block_id=0)
+    assert not actual_2  # Doesn't exist at block 0.