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

boa/environment.py

@@ -311,9 +311,18 @@ def _hook_trace_computation(self, computation, contract=None):
             child_contract = self._lookup_contract_fast(child.msg.code_address)
             self._hook_trace_computation(child, child_contract)
 
-    def get_code(self, address):
+    def get_code(self, address: _AddressType) -> bytes:
         return self.evm.get_code(Address(address))
 
+    def set_code(self, address: _AddressType, code: bytes) -> None:
+        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:
+        self.evm.set_storage(Address(address), slot, value)
+
     # function to time travel
     def time_travel(
         self,

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,12 @@ 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_balance(self, address, value):
+        raise NotImplementedError("Cannot use set_balance in network mode")
+
+    def set_code(self, address: _AddressType, code: bytes) -> None:
+        raise NotImplementedError("Cannot use set_code in network mode")
+
+    def set_storage(self, address: _AddressType, slot: int, value: int) -> None:
+        raise NotImplementedError("Cannot use set_storage in network mode")

boa/vm/py_evm.py

@@ -420,6 +420,15 @@ def set_balance(self, address: Address, value):
     def get_code(self, address: Address) -> bytes:
         return self.vm.state.get_code(address.canonical_address)
 
+    def set_code(self, address: Address, code: bytes) -> None:
+        self.vm.state.set_code(address.canonical_address, code)
+
+    def get_storage(self, address: Address, slot: int) -> int:
+        return self.vm.state.get_storage(address.canonical_address, slot)
+
+    def set_storage(self, address: Address, slot: int, value: int) -> None:
+        self.vm.state.set_storage(address.canonical_address, slot, value)
+
     def get_gas_limit(self):
         return self.vm.state.gas_limit
 

tests/unitary/test_fork_utils.py

@@ -0,0 +1,42 @@
+import pytest
+
+import boa
+from boa import NetworkEnv
+
+address = "0x0000000000000000000000000000000000000065"
+
+
+def test_set_balance_network():
+    with pytest.raises(NotImplementedError):
+        NetworkEnv.set_balance(boa.env, address, 0)
+
+
+def test_set_code_network():
+    with pytest.raises(NotImplementedError):
+        NetworkEnv.set_code(boa.env, address, b"")
+
+
+def test_set_storage_network():
+    with pytest.raises(NotImplementedError):
+        NetworkEnv.set_storage(boa.env, address, 0, 0)
+
+
+def test_balance():
+    assert boa.env.get_balance(address) == 0
+    balance = 1000
+    boa.env.set_balance(address, balance)
+    assert boa.env.get_balance(address) == balance
+
+
+def test_code():
+    assert boa.env.get_code(address) == b""
+    code = b"some test code"
+    boa.env.set_code(address, code)
+    assert boa.env.get_code(address) == code
+
+
+def test_storage():
+    storage = 12381920371289
+    assert boa.env.get_storage(address, 0) == 0
+    boa.env.set_storage(address, 0, storage)
+    assert boa.env.get_storage(address, 0) == storage