From 932f365ce6f9f00c8f89f513e608d0718288de0f Mon Sep 17 00:00:00 2001
From: spencer-tb <spencer@spencertaylorbrown.uk>
Date: Sun, 24 Mar 2024 14:20:18 +0800
Subject: [PATCH 1/4] feat(fw): add verkle sub-command, boiler plate for post
 vkt verify.

---
 .../spec/base/base_test.py                    | 16 +++++++----
 .../spec/blockchain/blockchain_test.py        |  8 ++++--
 src/evm_transition_tool/geth.py               | 28 +++++++++++++++++++
 src/evm_transition_tool/transition_tool.py    |  1 +
 4 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/src/ethereum_test_tools/spec/base/base_test.py b/src/ethereum_test_tools/spec/base/base_test.py
index 11df359e72..9ed7eb001b 100644
--- a/src/ethereum_test_tools/spec/base/base_test.py
+++ b/src/ethereum_test_tools/spec/base/base_test.py
@@ -44,15 +44,11 @@ def verify_transactions(txs: List[Transaction] | None, result) -> List[int]:
     return list(rejected_txs.keys())
 
 
-def verify_post_alloc(
-    *, expected_post: Mapping, got_alloc: Mapping, got_vkt: Optional[Mapping] = None
-):
+def verify_post_alloc(*, expected_post: Mapping, got_alloc: Mapping):
     """
-    Verify that an allocation matches the expected post in the test.
+    Verify that the final allocation mapping matches the expected post in the test.
     Raises exception on unexpected values.
     """
-    # TODO: If VKT is not None, we have an overlay of the VKT on top of the alloc and we need
-    # to verify that
     got_alloc_normalized: Dict[Address, Any] = {
         Address(address): got_alloc[address] for address in got_alloc
     }
@@ -69,6 +65,14 @@ def verify_post_alloc(
                     raise Exception(f"expected account not found: {address}")
 
 
+def verify_post_vkt(*, expected_post: Mapping, got_vkt: Mapping):
+    """
+    Verify that the final verkle tree mapping matches the expected post in the test.
+    Raises exception on unexpected values.
+    """
+    # TODO: Add the use the verkle subcommand to get the keys for each value in the tree.
+
+
 def verify_result(result: Mapping, env: Environment):
     """
     Verify that values in the t8n result match the expected values.
diff --git a/src/ethereum_test_tools/spec/blockchain/blockchain_test.py b/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
index 8b068e2e32..7f7226e92d 100644
--- a/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
+++ b/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
@@ -32,6 +32,7 @@
     BaseFixture,
     BaseTest,
     verify_post_alloc,
+    verify_post_vkt,
     verify_result,
     verify_transactions,
 )
@@ -314,10 +315,13 @@ def network_info(self, fork: Fork, eips: Optional[List[int]] = None):
 
     def verify_post_state(self, *, t8n, alloc, vkt=None):
         """
-        Verifies the post alloc after all block/s or payload/s are generated.
+        Verifies the post state after all block/s or payload/s are generated.
         """
         try:
-            verify_post_alloc(expected_post=self.post, got_alloc=alloc, got_vkt=vkt)
+            if vkt is not None and alloc is None:
+                verify_post_vkt(expected_post=self.post, got_vkt=vkt)
+            else:
+                verify_post_alloc(expected_post=self.post, got_alloc=alloc)
         except Exception as e:
             print_traces(t8n.get_traces())
             raise e
diff --git a/src/evm_transition_tool/geth.py b/src/evm_transition_tool/geth.py
index 0ad7efd32f..b1e464351b 100644
--- a/src/evm_transition_tool/geth.py
+++ b/src/evm_transition_tool/geth.py
@@ -13,6 +13,7 @@
 import pytest
 
 from ethereum_test_forks import Fork
+from ethereum_test_tools import Address
 
 from .transition_tool import FixtureFormats, TransitionTool, dump_files_to_directory
 
@@ -27,6 +28,7 @@ class GethTransitionTool(TransitionTool):
     t8n_subcommand: Optional[str] = "t8n"
     statetest_subcommand: Optional[str] = "statetest"
     blocktest_subcommand: Optional[str] = "blocktest"
+    verkle_subcommand: Optional[str] = "verkle"
 
     binary: Path
     cached_version: Optional[str] = None
@@ -125,3 +127,29 @@ def verify_fixture(
                 f"Failed to verify fixture via: '{' '.join(command)}'. "
                 f"Error: '{result.stderr.decode()}'"
             )
+
+        def verkle_tree_key(self, account: Address, storage_slot: Optional[str]) -> str:
+            """
+            Returns the verkle tree key for the input account using the verkle subcommand.
+            Optionally the key for the storage slot if specified.
+            """
+            command: list[str] = [str(self.binary)]
+            command.append(self.verkle_subcommand)
+
+            command.append(str(account))
+            if storage_slot:
+                command.append(storage_slot)
+
+            command = [str(self.binary), self.verkle_subcommand]
+            result = subprocess.run(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+
+            if result.returncode != 0:
+                raise Exception(
+                    f"Failed to run verkle subcommand: '{' '.join(command)}'. "
+                    f"Error: '{result.stderr.decode()}'"
+                )
+            return result.stdout.decode()
diff --git a/src/evm_transition_tool/transition_tool.py b/src/evm_transition_tool/transition_tool.py
index d6488b0184..5ce4f9b40e 100644
--- a/src/evm_transition_tool/transition_tool.py
+++ b/src/evm_transition_tool/transition_tool.py
@@ -102,6 +102,7 @@ class TransitionTool:
     blocktest_subcommand: Optional[str] = None
     cached_version: Optional[str] = None
     t8n_use_stream: bool = True
+    verkle_subcommand: Optional[str] = None
 
     # Abstract methods that each tool must implement
 

From d4943ae4e163613557935e4d0d2830b9e77b5b01 Mon Sep 17 00:00:00 2001
From: spencer-tb <spencer@spencertaylorbrown.uk>
Date: Sun, 24 Mar 2024 15:15:35 +0800
Subject: [PATCH 2/4] feat(fw): add initial post vkt verification checks.

---
 .../spec/base/base_test.py                    | 29 +++++++++++++++++--
 .../spec/blockchain/blockchain_test.py        |  2 +-
 src/evm_transition_tool/geth.py               | 11 ++++++-
 src/evm_transition_tool/transition_tool.py    | 12 +++++++-
 4 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/src/ethereum_test_tools/spec/base/base_test.py b/src/ethereum_test_tools/spec/base/base_test.py
index 9ed7eb001b..f8ea96e073 100644
--- a/src/ethereum_test_tools/spec/base/base_test.py
+++ b/src/ethereum_test_tools/spec/base/base_test.py
@@ -65,12 +65,37 @@ def verify_post_alloc(*, expected_post: Mapping, got_alloc: Mapping):
                     raise Exception(f"expected account not found: {address}")
 
 
-def verify_post_vkt(*, expected_post: Mapping, got_vkt: Mapping):
+def verify_post_vkt(transition_tool: TransitionTool, expected_post: Mapping, got_vkt: Mapping):
     """
     Verify that the final verkle tree mapping matches the expected post in the test.
     Raises exception on unexpected values.
     """
-    # TODO: Add the use the verkle subcommand to get the keys for each value in the tree.
+    # TODO: Add this check before filling within the filler cli plugin.
+    if not transition_tool.verkle_subcommand:
+        raise Exception("Only geth's evm tool is supported to verify verkle trees.")
+
+    # Convert the expected post alloc to a verkle tree for comparison.
+    # TODO: Maintain an intermediate type VerkleKeyMap: that maps verkle keys to alloc keys
+    expected_vkt = transition_tool.post_alloc_to_vkt(post_alloc=expected_post)
+
+    # Check for keys that are missing or unexpected in the actual VKT
+    missing_keys = [key for key in expected_vkt if key not in got_vkt]
+    unexpected_keys = [key for key in got_vkt if key not in expected_vkt]
+    if missing_keys or unexpected_keys:
+        error_messages = []
+        if missing_keys:
+            error_messages.append(f"Missing keys in actual VKT: {missing_keys}")
+        if unexpected_keys:
+            error_messages.append(f"Unexpected keys in actual VKT: {unexpected_keys}")
+        raise Exception("Verkle tree mismatch:\n" + "\n".join(error_messages))
+
+    # Compare the values for each key in the expected VKT
+    for key, expected_value in expected_vkt.items():
+        actual_value = got_vkt.get(key)
+        if expected_value != actual_value:
+            raise Exception(
+                f"Mismatch at key {key}: expected {expected_value}, got {actual_value}"
+            )
 
 
 def verify_result(result: Mapping, env: Environment):
diff --git a/src/ethereum_test_tools/spec/blockchain/blockchain_test.py b/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
index 7f7226e92d..9d334b80c0 100644
--- a/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
+++ b/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
@@ -319,7 +319,7 @@ def verify_post_state(self, *, t8n, alloc, vkt=None):
         """
         try:
             if vkt is not None and alloc is None:
-                verify_post_vkt(expected_post=self.post, got_vkt=vkt)
+                verify_post_vkt(transition_tool=t8n, expected_post=self.post, got_vkt=vkt)
             else:
                 verify_post_alloc(expected_post=self.post, got_alloc=alloc)
         except Exception as e:
diff --git a/src/evm_transition_tool/geth.py b/src/evm_transition_tool/geth.py
index b1e464351b..c86f426108 100644
--- a/src/evm_transition_tool/geth.py
+++ b/src/evm_transition_tool/geth.py
@@ -8,7 +8,7 @@
 import textwrap
 from pathlib import Path
 from re import compile
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import pytest
 
@@ -153,3 +153,12 @@ def verkle_tree_key(self, account: Address, storage_slot: Optional[str]) -> str:
                     f"Error: '{result.stderr.decode()}'"
                 )
             return result.stdout.decode()
+
+        def post_alloc_to_vkt(self, post_alloc: Dict[str, Dict[str, str]]) -> Dict[str, Any]:
+            """
+            Converts the expected post alloc to verkle tree representation using the verkle
+            subcommand.
+            """
+            raise Exception(
+                "`create_post_vkt()` function is not supported by this tool. Use geth's evm tool."
+            )
diff --git a/src/evm_transition_tool/transition_tool.py b/src/evm_transition_tool/transition_tool.py
index 5ce4f9b40e..9301d1071c 100644
--- a/src/evm_transition_tool/transition_tool.py
+++ b/src/evm_transition_tool/transition_tool.py
@@ -14,7 +14,7 @@
 from itertools import groupby
 from pathlib import Path
 from re import Pattern
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Type
 
 from ethereum_test_forks import Fork
 
@@ -647,3 +647,13 @@ def verify_fixture(
         raise Exception(
             "The `verify_fixture()` function is not supported by this tool. Use geth's evm tool."
         )
+
+    def post_alloc_to_vkt(self, post_alloc: Mapping) -> Mapping:
+        """
+        Converts the expected post alloc to verkle tree representation using the verkle subcommand.
+
+        Currently only implemented by geth's evm.
+        """
+        raise Exception(
+            "The `create_post_vkt()` function is not supported by this tool. Use geth's evm tool."
+        )

From da929b290eaec0b2681c099c09490f048bcc3b01 Mon Sep 17 00:00:00 2001
From: spencer-tb <spencer@spencertaylorbrown.uk>
Date: Mon, 25 Mar 2024 15:43:16 +0800
Subject: [PATCH 3/4] feat(fw): create alloc to vkt func, add test for it.

---
 src/evm_transition_tool/geth.py               | 92 ++++++++++++-------
 .../tests/test_alloc_to_vkt.py                | 48 ++++++++++
 2 files changed, 107 insertions(+), 33 deletions(-)
 create mode 100644 src/evm_transition_tool/tests/test_alloc_to_vkt.py

diff --git a/src/evm_transition_tool/geth.py b/src/evm_transition_tool/geth.py
index c86f426108..933ae0333e 100644
--- a/src/evm_transition_tool/geth.py
+++ b/src/evm_transition_tool/geth.py
@@ -8,12 +8,13 @@
 import textwrap
 from pathlib import Path
 from re import compile
-from typing import Any, Dict, Optional
+from typing import Mapping, Optional
 
 import pytest
+from ethereum.crypto.hash import keccak256
 
 from ethereum_test_forks import Fork
-from ethereum_test_tools import Address
+from ethereum_test_tools import Hash
 
 from .transition_tool import FixtureFormats, TransitionTool, dump_files_to_directory
 
@@ -128,37 +129,62 @@ def verify_fixture(
                 f"Error: '{result.stderr.decode()}'"
             )
 
-        def verkle_tree_key(self, account: Address, storage_slot: Optional[str]) -> str:
-            """
-            Returns the verkle tree key for the input account using the verkle subcommand.
-            Optionally the key for the storage slot if specified.
-            """
-            command: list[str] = [str(self.binary)]
-            command.append(self.verkle_subcommand)
-
-            command.append(str(account))
-            if storage_slot:
-                command.append(storage_slot)
-
-            command = [str(self.binary), self.verkle_subcommand]
-            result = subprocess.run(
-                command,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-            )
+    def verkle_tree_key(self, account: str, storage_slot: Optional[str] = None) -> str:
+        """
+        Returns the verkle tree key for the input account using the verkle subcommand.
+        Optionally the key for the storage slot if specified.
+        """
+        command = [
+            str(self.binary),
+            str(self.t8n_subcommand),
+            str(self.verkle_subcommand),
+            str(account),
+        ]
+        if storage_slot:
+            command.append(storage_slot)
+
+        result = subprocess.run(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
 
-            if result.returncode != 0:
-                raise Exception(
-                    f"Failed to run verkle subcommand: '{' '.join(command)}'. "
-                    f"Error: '{result.stderr.decode()}'"
-                )
-            return result.stdout.decode()
-
-        def post_alloc_to_vkt(self, post_alloc: Dict[str, Dict[str, str]]) -> Dict[str, Any]:
-            """
-            Converts the expected post alloc to verkle tree representation using the verkle
-            subcommand.
-            """
+        if result.returncode != 0:
             raise Exception(
-                "`create_post_vkt()` function is not supported by this tool. Use geth's evm tool."
+                f"Failed to run verkle subcommand: '{' '.join(command)}'. "
+                f"Error: '{result.stderr.decode()}'"
             )
+        return result.stdout.decode().strip()  # strip the newline character
+
+    def post_alloc_to_vkt(self, post_alloc: Mapping) -> Mapping:
+        """
+        Converts the expected post alloc to verkle tree representation using the verkle
+        subcommand.
+        """
+        vkt = {}
+        for address, account in post_alloc.items():
+            # Add the account address: value is simply "0x000...000"
+            address_key = self.verkle_tree_key(address)
+            vkt[address_key] = Hash(0)
+
+            # Add account balance: numbers are little-endian
+            balance_key = address_key[:-2] + "01"
+            balance_value = Hash(int(account["balance"], 16).to_bytes(32, "little"))
+            vkt[balance_key] = balance_value
+
+            # Add account nonce: numbers are little-endian
+            nonce_key = address_key[:-2] + "02"
+            nonce_value = Hash(int(account["nonce"], 16).to_bytes(32, "little"))
+            vkt[nonce_key] = nonce_value
+
+            # Add account code hash: keccak256 hash of the code
+            code_hash_key = address_key[:-2] + "03"
+            code_hash_value = Hash(keccak256(bytes.fromhex(account["code"][2:])).hex())
+            vkt[code_hash_key] = code_hash_value
+
+            # Add account storage: each slot has a unique key
+            for slot, value in account["storage"].items():
+                slot_key = self.verkle_tree_key(address, slot)
+                vkt[slot_key] = Hash(value)
+
+        return vkt
diff --git a/src/evm_transition_tool/tests/test_alloc_to_vkt.py b/src/evm_transition_tool/tests/test_alloc_to_vkt.py
new file mode 100644
index 0000000000..a33810c098
--- /dev/null
+++ b/src/evm_transition_tool/tests/test_alloc_to_vkt.py
@@ -0,0 +1,48 @@
+"""
+Test the verkle tree subcommand from the geth transition tool.
+"""
+
+import pytest
+
+from evm_transition_tool import GethTransitionTool
+
+
+@pytest.mark.parametrize(
+    "post_alloc, expected_vkt",
+    [
+        (
+            {
+                "0x0000000000000000000000000000000000000100": {
+                    "nonce": "0x01",
+                    "balance": "0x01",
+                    "code": "0x60203560003555",
+                    "storage": {"0x0a": "0x0b"},
+                },
+            },
+            {
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36300": "0x0000000000000000000000000000000000000000000000000000000000000000",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36301": "0x0100000000000000000000000000000000000000000000000000000000000000",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36302": "0x0100000000000000000000000000000000000000000000000000000000000000",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd36303": "0x159c5cfa7fab15702c72a4141ef31b26c42827bd74ffc5671d95033227e662e7",  # noqa: E501
+                "0x31b64bbd0b09c1d09afea606bebb70bce80a2909189e513c924a0871bbd3634a": "0x000000000000000000000000000000000000000000000000000000000000000b",  # noqa: E501
+            },
+        ),
+    ],
+)
+def test_post_alloc_to_vkt(post_alloc, expected_vkt):
+    """
+    Verifies that the `post_alloc_to_vkt` method of the `GethTransitionTool` class.
+    """
+    t8n = GethTransitionTool()
+    result_vkt = t8n.post_alloc_to_vkt(post_alloc)
+
+    assert set(result_vkt.keys()) == set(
+        expected_vkt.keys()
+    ), "Keys in created verkle tree do not match the expected keys."
+
+    for key, expected_value in expected_vkt.items():
+        assert key in result_vkt, f"Key {key} is missing in created verkle tree."
+        assert result_vkt[key] == expected_value, (
+            f"Value for {key} in the created verkle tree does not match expected. "
+            f"Expected {expected_value}, got {result_vkt[key]}."
+        )

From 2023092f0f952106590b4f62dd278cf12f818705 Mon Sep 17 00:00:00 2001
From: spencer-tb <spencer@spencertaylorbrown.uk>
Date: Mon, 25 Mar 2024 22:14:48 +0800
Subject: [PATCH 4/4] feat(fw): update vkt post state verifaction.

---
 .../spec/base/base_test.py                    | 20 +++++++++----------
 .../spec/blockchain/blockchain_test.py        |  6 ++++--
 src/evm_transition_tool/geth.py               | 14 ++++++-------
 .../tests/test_alloc_to_vkt.py                |  1 +
 4 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/src/ethereum_test_tools/spec/base/base_test.py b/src/ethereum_test_tools/spec/base/base_test.py
index f8ea96e073..8e7010af14 100644
--- a/src/ethereum_test_tools/spec/base/base_test.py
+++ b/src/ethereum_test_tools/spec/base/base_test.py
@@ -65,7 +65,9 @@ def verify_post_alloc(*, expected_post: Mapping, got_alloc: Mapping):
                     raise Exception(f"expected account not found: {address}")
 
 
-def verify_post_vkt(transition_tool: TransitionTool, expected_post: Mapping, got_vkt: Mapping):
+def verify_post_vkt(
+    transition_tool: TransitionTool, expected_post: Mapping, got_vkt: Mapping, got_alloc: Mapping
+):
     """
     Verify that the final verkle tree mapping matches the expected post in the test.
     Raises exception on unexpected values.
@@ -75,26 +77,22 @@ def verify_post_vkt(transition_tool: TransitionTool, expected_post: Mapping, got
         raise Exception("Only geth's evm tool is supported to verify verkle trees.")
 
     # Convert the expected post alloc to a verkle tree for comparison.
-    # TODO: Maintain an intermediate type VerkleKeyMap: that maps verkle keys to alloc keys
     expected_vkt = transition_tool.post_alloc_to_vkt(post_alloc=expected_post)
 
-    # Check for keys that are missing or unexpected in the actual VKT
+    # Check for keys that are missing the actual VKT
     missing_keys = [key for key in expected_vkt if key not in got_vkt]
+    if missing_keys:
+        raise Exception(f"Missing keys in actual VKT: {missing_keys}")
+
+    # TODO: how to determine what is unexpected, i.e TestAddress is expected but not in post state
     unexpected_keys = [key for key in got_vkt if key not in expected_vkt]
-    if missing_keys or unexpected_keys:
-        error_messages = []
-        if missing_keys:
-            error_messages.append(f"Missing keys in actual VKT: {missing_keys}")
-        if unexpected_keys:
-            error_messages.append(f"Unexpected keys in actual VKT: {unexpected_keys}")
-        raise Exception("Verkle tree mismatch:\n" + "\n".join(error_messages))
 
     # Compare the values for each key in the expected VKT
     for key, expected_value in expected_vkt.items():
         actual_value = got_vkt.get(key)
         if expected_value != actual_value:
             raise Exception(
-                f"Mismatch at key {key}: expected {expected_value}, got {actual_value}"
+                f"VKT mismatch at key {key}: expected {expected_value}, got {actual_value}"
             )
 
 
diff --git a/src/ethereum_test_tools/spec/blockchain/blockchain_test.py b/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
index 9d334b80c0..a62c6b6361 100644
--- a/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
+++ b/src/ethereum_test_tools/spec/blockchain/blockchain_test.py
@@ -318,8 +318,10 @@ def verify_post_state(self, *, t8n, alloc, vkt=None):
         Verifies the post state after all block/s or payload/s are generated.
         """
         try:
-            if vkt is not None and alloc is None:
-                verify_post_vkt(transition_tool=t8n, expected_post=self.post, got_vkt=vkt)
+            if vkt is not None:
+                verify_post_vkt(
+                    transition_tool=t8n, expected_post=self.post, got_vkt=vkt, got_alloc=alloc
+                )
             else:
                 verify_post_alloc(expected_post=self.post, got_alloc=alloc)
         except Exception as e:
diff --git a/src/evm_transition_tool/geth.py b/src/evm_transition_tool/geth.py
index 933ae0333e..16f6aeaa92 100644
--- a/src/evm_transition_tool/geth.py
+++ b/src/evm_transition_tool/geth.py
@@ -165,26 +165,26 @@ def post_alloc_to_vkt(self, post_alloc: Mapping) -> Mapping:
         for address, account in post_alloc.items():
             # Add the account address: value is simply "0x000...000"
             address_key = self.verkle_tree_key(address)
-            vkt[address_key] = Hash(0)
+            vkt[address_key] = Hash(0).hex()
 
             # Add account balance: numbers are little-endian
             balance_key = address_key[:-2] + "01"
-            balance_value = Hash(int(account["balance"], 16).to_bytes(32, "little"))
+            balance_value = Hash(account.balance.to_bytes(32, "little")).hex()
             vkt[balance_key] = balance_value
 
             # Add account nonce: numbers are little-endian
             nonce_key = address_key[:-2] + "02"
-            nonce_value = Hash(int(account["nonce"], 16).to_bytes(32, "little"))
+            nonce_value = Hash(account.nonce.to_bytes(32, "little")).hex()
             vkt[nonce_key] = nonce_value
 
             # Add account code hash: keccak256 hash of the code
             code_hash_key = address_key[:-2] + "03"
-            code_hash_value = Hash(keccak256(bytes.fromhex(account["code"][2:])).hex())
+            code_hash_value = Hash(keccak256(account.code)).hex()
             vkt[code_hash_key] = code_hash_value
 
             # Add account storage: each slot has a unique key
-            for slot, value in account["storage"].items():
-                slot_key = self.verkle_tree_key(address, slot)
-                vkt[slot_key] = Hash(value)
+            for slot, value in account.storage.data.items():
+                slot_key = self.verkle_tree_key(address, Hash(slot).hex())
+                vkt[slot_key] = Hash(value).hex()
 
         return vkt
diff --git a/src/evm_transition_tool/tests/test_alloc_to_vkt.py b/src/evm_transition_tool/tests/test_alloc_to_vkt.py
index a33810c098..05dd0a6098 100644
--- a/src/evm_transition_tool/tests/test_alloc_to_vkt.py
+++ b/src/evm_transition_tool/tests/test_alloc_to_vkt.py
@@ -7,6 +7,7 @@
 from evm_transition_tool import GethTransitionTool
 
 
+# TODO: Update to use correct types.
 @pytest.mark.parametrize(
     "post_alloc, expected_vkt",
     [