From afc2c2ed637eaba2b8d9a31b609929ec5ea854ac Mon Sep 17 00:00:00 2001
From: Jochem Brouwer <jochembrouwer96@gmail.com>
Date: Fri, 13 Sep 2024 21:13:24 +0200
Subject: [PATCH] new(test): EIP-7702 + EIP-1153: test that TransientStorage
 stays at correct address

---
 .../eip7702_set_code_tx/test_set_code_txs.py  | 101 ++++++++++++++++--
 1 file changed, 90 insertions(+), 11 deletions(-)

diff --git a/tests/prague/eip7702_set_code_tx/test_set_code_txs.py b/tests/prague/eip7702_set_code_tx/test_set_code_txs.py
index fea63a5685..fce89e8c0b 100644
--- a/tests/prague/eip7702_set_code_tx/test_set_code_txs.py
+++ b/tests/prague/eip7702_set_code_tx/test_set_code_txs.py
@@ -378,6 +378,81 @@ def test_set_code_to_tstore_reentry(
     )
 
 
+@pytest.mark.parametrize(
+    "call_opcode",
+    [
+        Op.CALL,
+        Op.STATICCALL,
+    ],
+)
+@pytest.mark.parametrize("call_eoa_first", [True, False])
+def test_set_code_to_tstore_available_at_correct_address(
+    state_test: StateTestFiller,
+    pre: Alloc,
+    call_opcode: Op,
+    call_eoa_first: bool,
+):
+    """
+    Test TLOADing from slot 2 and then SSTORE this in slot 1, then TSTORE 3 in slot 2.
+    This is done both from the EOA which is delegated to account A, and then A is called.
+    The storage should stay empty on both the EOA and the delegated account.
+    """
+    auth_signer = pre.fund_eoa(auth_account_start_balance)
+
+    storage_slot = 1
+    tload_slot = 2
+    tstore_value = 3
+
+    tstore_check_code = Op.SSTORE(storage_slot, Op.TLOAD(tload_slot)) + Op.TSTORE(
+        tload_slot, tstore_value
+    )
+
+    set_code_to_address = pre.deploy_contract(tstore_check_code)
+
+    def make_call(call_type: Op, call_eoa: bool) -> Bytecode:
+
+        call_target = auth_signer if call_eoa else set_code_to_address
+        if call_type == Op.STATICCALL:
+            return call_type(Op.GAS(), call_target, 0, 0, 0, 0)
+        else:
+            # call_type = Op.CALL
+            return call_type(Op.GAS(), call_target, 0, 0, 0, 0, 0)
+
+    chain_code = make_call(call_type=call_opcode, call_eoa=call_eoa_first) + make_call(
+        call_type=call_opcode, call_eoa=not call_eoa_first
+    )
+
+    target_call_chain_address = pre.deploy_contract(chain_code)
+
+    tx = Transaction(
+        gas_limit=100_000,
+        to=target_call_chain_address,
+        value=0,
+        authorization_list=[
+            AuthorizationTuple(
+                address=set_code_to_address,
+                nonce=0,
+                signer=auth_signer,
+            ),
+        ],
+        sender=pre.fund_eoa(),
+    )
+
+    state_test(
+        env=Environment(),
+        pre=pre,
+        tx=tx,
+        post={
+            auth_signer: Account(
+                storage={storage_slot: 0},
+            ),
+            set_code_to_address: Account(
+                storage={storage_slot: 0},
+            ),
+        },
+    )
+
+
 @pytest.mark.parametrize(
     "external_sendall_recipient",
     [False, True],
@@ -659,9 +734,11 @@ def test_set_code_call_set_code(
             auth_signer_1: Account(
                 nonce=1,
                 code=Spec.delegation_designation(set_code_to_address_1),
-                storage=storage_1
-                if call_opcode in [Op.CALL, Op.STATICCALL, Op.EXTCALL, Op.EXTSTATICCALL]
-                else storage_1 + storage_2,
+                storage=(
+                    storage_1
+                    if call_opcode in [Op.CALL, Op.STATICCALL, Op.EXTCALL, Op.EXTSTATICCALL]
+                    else storage_1 + storage_2
+                ),
                 balance=(0 if call_opcode in [Op.CALL, Op.EXTCALL] else value)
                 + auth_account_start_balance,
             ),
@@ -974,9 +1051,11 @@ def test_ext_code_on_set_code(
         pre=pre,
         tx=tx,
         post={
-            set_code_to_address: Account.NONEXISTENT
-            if set_code_type == AddressType.EMPTY_ACCOUNT
-            else Account(storage={}),
+            set_code_to_address: (
+                Account.NONEXISTENT
+                if set_code_type == AddressType.EMPTY_ACCOUNT
+                else Account(storage={})
+            ),
             auth_signer: Account(
                 nonce=1,
                 code=Spec.delegation_designation(set_code_to_address),
@@ -1866,11 +1945,11 @@ def test_set_code_invalid_authorization_tuple(
         authorization_list = [
             AuthorizationTuple(
                 address=set_code_to_address,
-                nonce=1
-                if invalidity_reason == InvalidityReason.NONCE
-                else [0, 1]
-                if invalidity_reason == InvalidityReason.MULTIPLE_NONCE
-                else 0,
+                nonce=(
+                    1
+                    if invalidity_reason == InvalidityReason.NONCE
+                    else [0, 1] if invalidity_reason == InvalidityReason.MULTIPLE_NONCE else 0
+                ),
                 chain_id=2 if invalidity_reason == InvalidityReason.CHAIN_ID else 0,
                 signer=auth_signer,
             )