fix test; make storage mock work on iterable values

kkrt-labs · Nov 14, 2024 · 3fb7fed · 3fb7fed
1 parent fe28e81
commit 3fb7fed
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 43 deletions.
diff --git a/cairo_zero/kakarot/library.cairo b/cairo_zero/kakarot/library.cairo
@@ -179,7 +179,7 @@ namespace Kakarot {
             return ();
         }
 
-        Kakarot_base_fee.write('current_block', (base_fee, starting_block));
+        Kakarot_base_fee.write('current_block', (next_base_fee, starting_block));
         return ();
     }
 

diff --git a/cairo_zero/tests/src/kakarot/test_kakarot.py b/cairo_zero/tests/src/kakarot/test_kakarot.py
@@ -165,14 +165,7 @@ def test_set_base_fee_should_set_next_block_fee(self, cairo_run):
             get_storage_var_address(
                 "Kakarot_base_fee", int.from_bytes(b"next_block", "big")
             ),
-            value=1,
-        )
-        @SyscallHandler.patch(
-            get_storage_var_address(
-                "Kakarot_base_fee", int.from_bytes(b"next_block", "big")
-            )
-            + 1,
-            value=0x101,
+            value=[1, 0x101],
         )
         @patch.object(SyscallHandler, "block_number", 0x102)
         def test_set_base_fee_should_overwrite_current_block_fee_if_next_block_is_applicable(
@@ -213,27 +206,13 @@ def test_set_base_fee_should_overwrite_current_block_fee_if_next_block_is_applic
             get_storage_var_address(
                 "Kakarot_base_fee", int.from_bytes(b"current_block", "big")
             ),
-            value=1,
-        )
-        @SyscallHandler.patch(
-            get_storage_var_address(
-                "Kakarot_base_fee", int.from_bytes(b"current_block", "big")
-            )
-            + 1,
-            value=0x100,
+            value=[1, 0x100],
         )
         @SyscallHandler.patch(
             get_storage_var_address(
                 "Kakarot_base_fee", int.from_bytes(b"next_block", "big")
             ),
-            value=2,
-        )
-        @SyscallHandler.patch(
-            get_storage_var_address(
-                "Kakarot_base_fee", int.from_bytes(b"next_block", "big")
-            )
-            + 1,
-            value=0x101,
+            value=[2, 0x101],
         )
         @patch.object(SyscallHandler, "block_number", 0x100)
         def test_get_base_fee_should_return_current_block_fee_if_next_block_is_not_applicable(
@@ -247,27 +226,13 @@ def test_get_base_fee_should_return_current_block_fee_if_next_block_is_not_appli
             get_storage_var_address(
                 "Kakarot_base_fee", int.from_bytes(b"current_block", "big")
             ),
-            value=1,
-        )
-        @SyscallHandler.patch(
-            get_storage_var_address(
-                "Kakarot_base_fee", int.from_bytes(b"current_block", "big")
-            )
-            + 1,
-            value=0x100,
+            value=[1, 0x100],
         )
         @SyscallHandler.patch(
             get_storage_var_address(
                 "Kakarot_base_fee", int.from_bytes(b"next_block", "big")
             ),
-            value=2,
-        )
-        @SyscallHandler.patch(
-            get_storage_var_address(
-                "Kakarot_base_fee", int.from_bytes(b"next_block", "big")
-            )
-            + 1,
-            value=0x101,
+            value=[2, 0x101],
         )
         @patch.object(SyscallHandler, "block_number", 0x101)
         def test_get_base_fee_should_return_next_block_fee_if_applicable_and_update_current_block(

diff --git a/tests/utils/syscall_handler.py b/tests/utils/syscall_handler.py
@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from hashlib import sha256
-from typing import Optional, Union
+from typing import Iterable, Optional, Union
 from unittest import mock
 
 import ecdsa
@@ -609,7 +609,13 @@ def patch(
                 selector_if_storage = get_storage_var_address(target, *args)
             else:
                 selector_if_storage = target
-            cls.patches[selector_if_storage] = value
+
+            if isinstance(value, Iterable):
+                for i, v in enumerate(value):
+                    cls.patches[selector_if_storage + i] = v
+            else:
+                cls.patches[selector_if_storage] = value
+
         except AssertionError:
             pass