feat: add set_storage [APE-1093] (#55)

ape_foundry/provider.py

@@ -48,6 +48,7 @@
 from ape_foundry.constants import EVM_VERSION_BY_NETWORK
 
 from .exceptions import FoundryNotInstalledError, FoundryProviderError, FoundrySubprocessError
+from .utils import to_bytes32
 
 EPHEMERAL_PORTS_START = 49152
 EPHEMERAL_PORTS_END = 60999
@@ -698,6 +699,12 @@ def set_code(self, address: AddressType, code: ContractCode) -> bool:
         self._make_request("anvil_setCode", [address, code])
         return True
 
+    def set_storage(self, address: AddressType, slot: int, value: HexBytes):
+        self._make_request(
+            "anvil_setStorageAt",
+            [address, to_bytes32(slot).hex(), to_bytes32(value).hex()],
+        )
+
     def _eth_call(self, arguments: List) -> bytes:
         # Override from Web3Provider because foundry is pickier.
 

ape_foundry/utils.py

@@ -0,0 +1,24 @@
+from typing import Union
+
+from eth_utils import to_bytes, to_hex
+from ethpm_types import HexBytes
+
+
+# TODO: Upstream to ape core
+def to_bytes32(value: Union[int, str, bytes, HexBytes]) -> HexBytes:
+    if isinstance(value, int):
+        value = to_bytes(value)
+
+    elif isinstance(value, str):
+        if set(value.lower().replace("0x", "")) > set("0123456789abcdef"):
+            raise TypeError(f"'{value}' not valid hexstr")
+
+        value = to_bytes(hexstr=value)
+
+    elif not isinstance(value, bytes):
+        raise TypeError(f"Cannot convert type '{type(value)}' to 'bytes'")
+
+    if len(value) > 32:
+        raise ValueError(f"Value '{to_hex(value)}' must be <= 32 bytes")
+
+    return HexBytes(value.rjust(32, b"\x00"))

tests/test_provider.py

@@ -6,6 +6,7 @@
 from ape.exceptions import ContractLogicError
 from ape.pytest.contextmanagers import RevertsContextManager as reverts
 from ape.types import CallTreeNode, TraceFrame
+from eth_utils import to_int
 from evm_trace import CallType
 from hexbytes import HexBytes
 
@@ -199,6 +200,13 @@ def test_set_code(connected_provider, contract_container, owner):
     assert provider.get_code(contract.address) == code
 
 
+def test_set_storage(connected_provider, contract_container, owner):
+    contract = contract_container.deploy(sender=owner)
+    assert to_int(connected_provider.get_storage_at(contract.address, "0x2b5e3af16b1880000")) == 0
+    connected_provider.set_storage(contract.address, "0x2b5e3af16b1880000", "0x1")
+    assert to_int(connected_provider.get_storage_at(contract.address, "0x2b5e3af16b1880000")) == 1
+
+
 def test_return_value(connected_provider, contract_instance, owner):
     receipt = contract_instance.setAddress(owner.address, sender=owner)
     assert receipt.return_value == 123