From 3fb7fed47b91e3b0dbc083ef63e738723811d5a2 Mon Sep 17 00:00:00 2001
From: enitrat <msaug@protonmail.com>
Date: Thu, 14 Nov 2024 16:42:38 +0700
Subject: [PATCH] fix test; make storage mock work on iterable values

---
 cairo_zero/kakarot/library.cairo             |  2 +-
 cairo_zero/tests/src/kakarot/test_kakarot.py | 45 +++-----------------
 tests/utils/syscall_handler.py               | 10 ++++-
 3 files changed, 14 insertions(+), 43 deletions(-)

diff --git a/cairo_zero/kakarot/library.cairo b/cairo_zero/kakarot/library.cairo
index b6bb00fb2..48f1046f5 100644
--- 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
index 284c31a1d..5d1f0e2e6 100644
--- 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
index b94424c2b..5a70e0515 100644
--- 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