From 07f4d3099f9bf71e46414760d0be4696ae918546 Mon Sep 17 00:00:00 2001
From: El De-dog-lo <3859395+fubuloubu@users.noreply.github.com>
Date: Thu, 22 Jun 2023 18:17:34 -0400
Subject: [PATCH] feat: add `set_storage` [APE-1093] (#55)

---
 ape_foundry/provider.py |  7 +++++++
 ape_foundry/utils.py    | 24 ++++++++++++++++++++++++
 tests/test_provider.py  |  8 ++++++++
 3 files changed, 39 insertions(+)
 create mode 100644 ape_foundry/utils.py

diff --git a/ape_foundry/provider.py b/ape_foundry/provider.py
index 55b9c97..5b210c0 100644
--- a/ape_foundry/provider.py
+++ b/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.
 
diff --git a/ape_foundry/utils.py b/ape_foundry/utils.py
new file mode 100644
index 0000000..5ccbc8b
--- /dev/null
+++ b/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"))
diff --git a/tests/test_provider.py b/tests/test_provider.py
index 47b96e3..3dfc044 100644
--- a/tests/test_provider.py
+++ b/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