From 8c8d4e8e7e8067be1cf01a171c630a92eba33562 Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Tue, 19 Mar 2024 08:47:16 -0400
Subject: [PATCH] adjust storage layout tests for cancun

---
 .../cli/storage_layout/test_storage_layout.py | 39 ++++++++++++++++---
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/tests/unit/cli/storage_layout/test_storage_layout.py b/tests/unit/cli/storage_layout/test_storage_layout.py
index 9724dd723c..ece2743b81 100644
--- a/tests/unit/cli/storage_layout/test_storage_layout.py
+++ b/tests/unit/cli/storage_layout/test_storage_layout.py
@@ -1,4 +1,21 @@
 from vyper.compiler import compile_code
+from vyper.evm.opcodes import version_check
+
+
+def _adjust_storage_layout_for_cancun(layout):
+    def _go(layout):
+        for _varname, item in layout.items():
+            if "slot" in item and isinstance(item["slot"], int):
+                item["slot"] -= 1
+            else:
+                # recurse to submodule
+                _go(item)
+
+    if version_check(begin="cancun"):
+        layout["transient_storage_layout"] = {
+            "$.nonreentrant_key": layout["storage_layout"].pop("$.nonreentrant_key")
+        }
+        _go(layout["storage_layout"])
 
 
 def test_storage_layout():
@@ -40,13 +57,18 @@ def public_foo3():
 
     out = compile_code(code, output_formats=["layout"])
 
-    assert out["layout"]["storage_layout"] == {
-        "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
-        "foo": {"slot": 1, "type": "HashMap[address, uint256]"},
-        "arr": {"slot": 2, "type": "DynArray[uint256, 3]"},
-        "baz": {"slot": 6, "type": "Bytes[65]"},
-        "bar": {"slot": 10, "type": "uint256"},
+    expected = {
+        "storage_layout": {
+            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
+            "foo": {"slot": 1, "type": "HashMap[address, uint256]"},
+            "arr": {"slot": 2, "type": "DynArray[uint256, 3]"},
+            "baz": {"slot": 6, "type": "Bytes[65]"},
+            "bar": {"slot": 10, "type": "uint256"},
+        }
     }
+    _adjust_storage_layout_for_cancun(expected)
+
+    assert out["layout"] == expected
 
 
 def test_storage_and_immutables_layout():
@@ -71,6 +93,7 @@ def __init__():
             "name": {"slot": 1, "type": "String[32]"},
         },
     }
+    _adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -120,6 +143,7 @@ def __init__():
             "a_library": {"supply": {"slot": 3, "type": "uint256"}},
         },
     }
+    _adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -169,6 +193,7 @@ def __init__():
             "counter2": {"slot": 3, "type": "uint256"},
         },
     }
+    _adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -253,6 +278,7 @@ def bar():
             "a_library": {"supply": {"slot": 4, "type": "uint256"}},
         },
     }
+    _adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -334,6 +360,7 @@ def foo() -> uint256:
             "counter2": {"slot": 4, "type": "uint256"},
         },
     }
+    _adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout