feat: test_balance exhaustive checks example.

src/ethereum_test_specs/blockchain.py

@@ -466,6 +466,7 @@ def generate_block_data(
                     )
                 self.verify_witness(
                     t8n=t8n,
+                    fork=fork,
                     state_diff=transition_tool_output.result.state_diff,
                     witness_check=block.witness_check,
                 )
@@ -619,6 +620,7 @@ def verify_post_state(
     def verify_witness(
         self,
         t8n: TransitionTool,
+        fork: Fork,
         state_diff: StateDiff,
         witness_check: WitnessCheck,
     ) -> None:
@@ -629,13 +631,19 @@ def verify_witness(
         witness_check_state_diff, witness_check_address_mapping = t8n.get_witness_check_mapping(
             witness_check
         )
+        print("\nExpected witness check state diff:")
+        print(witness_check_state_diff.model_dump_json(indent=4))
+        system_contract_accounts = fork.pre_allocation_blockchain().items()
+        addresses_to_skip = [Address(address) for address, _ in system_contract_accounts]
+
         for stem_state_diff in state_diff.root:
             actual_stem = stem_state_diff.stem
             address = witness_check_address_mapping.get(actual_stem, None)
-            print("Stem we are checking from actual witness:")
-            print(str(actual_stem))
-            print("Address we are checking from actual witness:")
-            print(str(address))
+            print(f"\nChecking witness for stem: {actual_stem} at address: {address}")
+            # TODO: skip system contract addresses for now
+            if address in addresses_to_skip:
+                print(f"Skipping system contract address in witness check: {address}")
+                continue
             # check for stem in the expected witness check
             expected_stem_state_diff = next(
                 (sd for sd in witness_check_state_diff.root if sd.stem == actual_stem), None
@@ -679,7 +687,9 @@ def verify_witness(
                         )
                     )
                 # check the current value of the actual suffix state diff matches the expected
-                if actual_current_value != expected_suffix_state_diff.current_value:
+                if str(actual_current_value) != str(
+                    expected_suffix_state_diff.current_value
+                ):  # TODO: temp fix str casting
                     raise ValueError(
                         "Witness check failed - current value mismatch. The stem and suffix"
                         " exist.\n\n"

src/ethereum_test_types/verkle/types.py

@@ -10,6 +10,7 @@
 from pydantic.functional_serializers import model_serializer
 
 from ethereum_test_base_types import Address, CamelModel, HexNumber, PaddedFixedSizeBytes
+from ethereum_test_forks import Fork, Verkle
 from ethereum_test_types import Account
 
 IPA_PROOF_DEPTH = 8
@@ -171,16 +172,22 @@ class AccountHeaderEntry(Enum):
         BASIC_DATA = 0
         CODEHASH = 1
 
-    def __init__(self) -> None:
+    def __init__(self, fork: Fork) -> None:
         """
         Initializes a WitnessCheck instance.
         """
+        assert fork >= Verkle, "WitnessCheck is only supported for Verkle fork and later"
         self.account_entries: List[
             Tuple[Address, WitnessCheck.AccountHeaderEntry, Optional[Hash]]
         ] = []
         self.storage_slots: List[Tuple[Address, int, Optional[Hash]]] = []
         self.code_chunks: List[Tuple[Address, int, Optional[Hash]]] = []
 
+        # Add the pre-allocation accounts by default
+        pre_allocations = fork.pre_allocation_blockchain()
+        for address, account_data in pre_allocations.items():
+            self.add_account_full(Address(address), Account(**account_data))
+
     def __repr__(self) -> str:
         """
         Provides a detailed string representation of the WitnessCheck object for debugging.
@@ -197,13 +204,17 @@ def add_account_full(self, address: Address, account: Account | None) -> None:
         """
         Adds the address, nonce, balance, and code. Delays actual key computation until later.
         """
-        if account and account.code:
-            code_hash = Hash(keccak256(account.code))
-        else:
-            # keccak256 of empty byte array
-            code_hash = Hash(0xC5D2460186F7233C927E7DB2DCC703C0E500B653CA82273B7BFAD8045D85A470)
         self.add_account_basic_data(address, account)
-        self.add_account_codehash(address, code_hash)
+        if account:
+            if account.code:
+                code_hash = Hash(keccak256(account.code))
+            else:  # keccak256 of empty byte array
+                code_hash = Hash(
+                    0xC5D2460186F7233C927E7DB2DCC703C0E500B653CA82273B7BFAD8045D85A470
+                )
+            self.add_account_codehash(address, code_hash)
+        else:
+            self.add_account_codehash(address, None)
 
     def add_account_basic_data(self, address: Address, account: Account | None) -> None:
         """