feat[fork]: add set_code, get_storage, set_storage to boa.env

boa/environment.py

@@ -315,14 +315,12 @@ def get_code(self, address: _AddressType) -> bytes:
         return self.evm.get_code(Address(address))
 
     def set_code(self, address: _AddressType, code: bytes) -> None:
-        assert self.evm.is_forked, "The EVM is not forked, cannot set code"
         self.evm.set_code(Address(address), code)
 
     def get_storage(self, address: _AddressType, slot: int) -> int:
         return self.evm.get_storage(Address(address), slot)
 
     def set_storage(self, address: _AddressType, slot: int, value: int) -> None:
-        assert self.evm.is_forked, "The EVM is not forked, cannot set storage"
         self.evm.set_storage(Address(address), slot, value)
 
     # function to time travel

boa/network.py

@@ -8,7 +8,7 @@
 from eth_account import Account
 from requests.exceptions import HTTPError
 
-from boa.environment import Env
+from boa.environment import Env, _AddressType
 from boa.rpc import (
     RPC,
     EthereumRPC,
@@ -495,3 +495,11 @@ def _send_txn(self, from_, to=None, gas=None, value=None, data=None):
 
         t_obj = TraceObject(trace) if trace is not None else None
         return receipt, t_obj
+
+    def set_code(self, address: _AddressType, code: bytes) -> None:
+        assert self.evm.is_forked, "The EVM is not forked, cannot set code"
+        return super().set_code(address, code)
+
+    def set_storage(self, address: _AddressType, slot: int, value: int) -> None:
+        assert self.evm.is_forked, "The EVM is not forked, cannot set storage"
+        super().set_storage(address, slot, value)

boa/vm/fork.py

@@ -120,7 +120,10 @@ def fetch_multi(self, payload):
             # fetch_multi is called only with the missing payloads
             # map the results back to the original indices
             for result_ix, rpc_result in enumerate(self._rpc.fetch_multi(batch)):
-                key, item_ix = keys[result_ix]
+                try:
+                    key, item_ix = keys[result_ix]
+                except IndexError:
+                    continue
                 ret[item_ix] = rpc_result
                 self._db[key] = json.dumps(rpc_result).encode("utf-8")
 

tests/unitary/test_fork_utils.py

@@ -1,23 +1,22 @@
 import pytest
 
 import boa
-from boa.vm.fork import AccountDBFork
+from boa import NetworkEnv
 
 address = "0x0000000000000000000000000000000000000065"
 
 
 def test_code_non_fork():
     with pytest.raises(AssertionError):
-        boa.env.set_code(address, b"")
+        NetworkEnv.set_code(boa.env, address, b"")
 
 
 def test_storage_non_fork():
     with pytest.raises(AssertionError):
-        boa.env.set_storage(address, 0, 0)
+        NetworkEnv.set_storage(boa.env, address, 0, 0)
 
 
 def test_code():
-    boa.env.evm._set_account_db_class(AccountDBFork)
     assert boa.env.get_code(address) == b""
     code = b"some test code"
     boa.env.set_code(address, code)
@@ -26,7 +25,6 @@ def test_code():
 
 def test_storage(monkeypatch):
     storage = 12381920371289
-    boa.env.evm._set_account_db_class(AccountDBFork)
     assert boa.env.get_storage(address, 0) == 0
     boa.env.set_storage(address, 0, storage)
     assert boa.env.get_storage(address, 0) == storage