feat: set code

eth_tester/backends/base.py

@@ -81,6 +81,10 @@ def get_balance(self, account, block_number=None):
     def get_code(self, account, block_number=None):
         raise NotImplementedError("Must be implemented by subclasses")
 
+    @abstractmethod
+    def set_code(self, account, code):
+        raise NotImplementedError("Must be implemented by subclasses")
+
     #
     # Transactions
     #

eth_tester/backends/mock/main.py

@@ -284,6 +284,12 @@ def get_code(self, account, block_number=None):
         except KeyError:
             return b""
 
+    def set_code(self, account, code):
+        if account in self.account_state_lookup:
+            self.account_state_lookup[account]["code"] = code
+        else:
+            self.account_state_lookup[account] = {"code": code}
+
     #
     # Transactions
     #

eth_tester/backends/pyevm/main.py

@@ -538,6 +538,10 @@ def get_code(self, account, block_number="latest"):
         vm = _get_vm_for_block_number(self.chain, block_number)
         return vm.state.get_code(account)
 
+    def set_code(self, account, code):
+        vm = _get_vm_for_block_number(self.chain, block_number)
+        return vm.state.set_code(account, code)
+
     def get_base_fee(self, block_number="latest"):
         vm = _get_vm_for_block_number(self.chain, block_number)
         return vm.state.base_fee

eth_tester/main.py

@@ -254,6 +254,11 @@ def get_code(self, account, block_number="latest"):
         code = self.normalizer.normalize_outbound_code(raw_code)
         return code
 
+    def set_code(self, account, code):
+        self.validator.validate_inbound_account(account)
+        raw_account = self.normalizer.normalize_inbound_account(account)
+        self.backend.set_code(raw_account, code)
+
     def get_nonce(self, account, block_number="latest"):
         self.validator.validate_inbound_account(account)
         self.validator.validate_inbound_block_number(block_number)

eth_tester/utils/backend_testing.py

@@ -231,6 +231,11 @@ def test_get_code_account_without_code(self, eth_tester):
         code = eth_tester.get_code(BURN_ADDRESS)
         assert code == "0x"
 
+    def test_set_code(self, eth_tester):
+        eth_tester.set_code(BURN_ADDRESS, "0x123")
+        code = eth_tester.get_code(BURN_ADDRESS)
+        assert code == "0x123"
+
     def test_get_nonce(self, eth_tester):
         for account in eth_tester.get_accounts():
             nonce = eth_tester.get_nonce(account)