From 02ae59f60b9cb548dc36f7fb41f5f59eec7723c1 Mon Sep 17 00:00:00 2001
From: enitrat <msaug@protonmail.com>
Date: Tue, 19 Nov 2024 22:25:01 +0800
Subject: [PATCH] fix tests

---
 .../kakarot/precompiles/test_precompiles.py   | 35 +++++--------------
 cairo_zero/tests/src/kakarot/test_state.py    |  9 +++--
 tests/utils/constants.py                      | 12 ++++---
 3 files changed, 20 insertions(+), 36 deletions(-)

diff --git a/cairo_zero/tests/src/kakarot/precompiles/test_precompiles.py b/cairo_zero/tests/src/kakarot/precompiles/test_precompiles.py
index eaec65405..61dd90fc5 100644
--- a/cairo_zero/tests/src/kakarot/precompiles/test_precompiles.py
+++ b/cairo_zero/tests/src/kakarot/precompiles/test_precompiles.py
@@ -5,11 +5,9 @@
 from tests.utils.constants import (
     CAIRO_MESSAGE_GAS,
     CAIRO_PRECOMPILE_GAS,
-    FIRST_KAKAROT_PRECOMPILE_ADDRESS,
-    FIRST_ROLLUP_PRECOMPILE_ADDRESS,
-    LAST_ETHEREUM_PRECOMPILE_ADDRESS,
-    LAST_KAKAROT_PRECOMPILE_ADDRESS,
-    LAST_ROLLUP_PRECOMPILE_ADDRESS,
+    ETHEREUM_PRECOMPILES,
+    KAKAROT_PRECOMPILES,
+    ROLLUP_PRECOMPILES,
 )
 from tests.utils.syscall_handler import SyscallHandler
 
@@ -282,38 +280,23 @@ def test__cairo_message(
                 assert gas_used == CAIRO_MESSAGE_GAS
 
     class TestIsPrecompile:
-        @pytest.mark.parametrize(
-            "address", range(0, LAST_ETHEREUM_PRECOMPILE_ADDRESS + 2)
-        )
+        @pytest.mark.parametrize("address", range(0, ETHEREUM_PRECOMPILES[-1] + 2))
         def test__is_precompile_ethereum_precompiles(self, cairo_run, address):
             result = cairo_run("test__is_precompile", address=address)
-            assert result == (address in range(1, LAST_ETHEREUM_PRECOMPILE_ADDRESS + 1))
+            assert result == (address in ETHEREUM_PRECOMPILES)
 
         @pytest.mark.parametrize(
             "address",
-            range(FIRST_ROLLUP_PRECOMPILE_ADDRESS, LAST_ROLLUP_PRECOMPILE_ADDRESS + 2),
+            range(ROLLUP_PRECOMPILES[0], ROLLUP_PRECOMPILES[-1] + 2),
         )
         def test__is_precompile_rollup_precompiles(self, cairo_run, address):
             result = cairo_run("test__is_precompile", address=address)
-            assert result == (
-                address
-                in range(
-                    FIRST_ROLLUP_PRECOMPILE_ADDRESS, LAST_ROLLUP_PRECOMPILE_ADDRESS + 1
-                )
-            )
+            assert result == (address in ROLLUP_PRECOMPILES)
 
         @pytest.mark.parametrize(
             "address",
-            range(
-                FIRST_KAKAROT_PRECOMPILE_ADDRESS, LAST_KAKAROT_PRECOMPILE_ADDRESS + 2
-            ),
+            range(KAKAROT_PRECOMPILES[0], KAKAROT_PRECOMPILES[-1] + 2),
         )
         def test__is_precompile_kakarot_precompiles(self, cairo_run, address):
             result = cairo_run("test__is_precompile", address=address)
-            assert result == (
-                address
-                in range(
-                    FIRST_KAKAROT_PRECOMPILE_ADDRESS,
-                    LAST_KAKAROT_PRECOMPILE_ADDRESS + 1,
-                )
-            )
+            assert result == (address in KAKAROT_PRECOMPILES)
diff --git a/cairo_zero/tests/src/kakarot/test_state.py b/cairo_zero/tests/src/kakarot/test_state.py
index fd1e43b10..9b5cce632 100644
--- a/cairo_zero/tests/src/kakarot/test_state.py
+++ b/cairo_zero/tests/src/kakarot/test_state.py
@@ -4,9 +4,8 @@
     TX_ACCESS_LIST_ADDRESS_COST,
     TX_ACCESS_LIST_STORAGE_KEY_COST,
 )
-from web3 import Web3
 
-from tests.utils.constants import TRANSACTIONS
+from tests.utils.constants import ALL_PRECOMPILES, TRANSACTIONS
 from tests.utils.helpers import flatten_tx_access_list, merge_access_list
 from tests.utils.syscall_handler import SyscallHandler
 
@@ -89,9 +88,9 @@ class TestCachePreaccessedAddresses:
         @SyscallHandler.patch("IERC20.balanceOf", lambda *_: [0, 1])
         def test_should_cache_precompiles(self, cairo_run):
             state = cairo_run("test__cache_precompiles")
-            assert list(map(Web3.to_checksum_address, state["accounts"].keys())) == [
-                Web3.to_checksum_address(f"0x{i:040x}") for i in range(1, 11)
-            ]
+            assert [
+                int(address, 16) for address in state["accounts"].keys()
+            ] == ALL_PRECOMPILES
 
         @SyscallHandler.patch("IERC20.balanceOf", lambda *_: [0, 1])
         @pytest.mark.parametrize("transaction", TRANSACTIONS)
diff --git a/tests/utils/constants.py b/tests/utils/constants.py
index f97cd0022..34a88c922 100644
--- a/tests/utils/constants.py
+++ b/tests/utils/constants.py
@@ -2,6 +2,7 @@
 from time import time
 
 import pytest
+from ethereum.cancun.vm.precompiled_contracts.mapping import PRE_COMPILED_CONTRACTS
 
 from kakarot_scripts.constants import BLOCK_GAS_LIMIT
 
@@ -26,11 +27,12 @@
 TRANSACTION_GAS_LIMIT = BLOCK_GAS_LIMIT
 
 # PRECOMPILES
-LAST_ETHEREUM_PRECOMPILE_ADDRESS = 0x0A
-FIRST_ROLLUP_PRECOMPILE_ADDRESS = 0x100
-LAST_ROLLUP_PRECOMPILE_ADDRESS = 0x100
-FIRST_KAKAROT_PRECOMPILE_ADDRESS = 0x75001
-LAST_KAKAROT_PRECOMPILE_ADDRESS = 0x75004
+ETHEREUM_PRECOMPILES = [
+    int.from_bytes(address, "big") for address in PRE_COMPILED_CONTRACTS.keys()
+]
+ROLLUP_PRECOMPILES = [0x100]
+KAKAROT_PRECOMPILES = [0x75001, 0x75002, 0x75003, 0x75004]
+ALL_PRECOMPILES = [*ETHEREUM_PRECOMPILES, *ROLLUP_PRECOMPILES, *KAKAROT_PRECOMPILES]
 
 CAIRO_PRECOMPILE_GAS = 10000
 CAIRO_MESSAGE_GAS = 5000