From e1fb88a185b528ac7aaa2f0fcd5f4609911e38cd Mon Sep 17 00:00:00 2001
From: Mario Vega <marioevz@gmail.com>
Date: Tue, 24 Sep 2024 23:10:12 +0000
Subject: [PATCH] fix(fw): pydantic fixture format

---
 src/cli/gen_index.py                                 |  2 +-
 src/ethereum_test_fixtures/consume.py                | 11 ++++++++---
 src/pytest_plugins/consume/direct/test_via_direct.py |  6 +++---
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/cli/gen_index.py b/src/cli/gen_index.py
index 43441d26a1..8178e4d1c3 100644
--- a/src/cli/gen_index.py
+++ b/src/cli/gen_index.py
@@ -195,7 +195,7 @@ def generate_fixtures_index(
                         json_path=relative_file_path,
                         fixture_hash=fixture.info.get("hash", None),
                         fork=fixture.get_fork(),
-                        format=fixture.fixture_format_name,
+                        format=fixture.__class__,
                     )
                 )
 
diff --git a/src/ethereum_test_fixtures/consume.py b/src/ethereum_test_fixtures/consume.py
index 73639911e8..ab46ee9e1f 100644
--- a/src/ethereum_test_fixtures/consume.py
+++ b/src/ethereum_test_fixtures/consume.py
@@ -5,11 +5,12 @@
 import datetime
 import json
 from pathlib import Path
-from typing import List, TextIO
+from typing import Annotated, List, TextIO
 
-from pydantic import BaseModel, RootModel
+from pydantic import BaseModel, PlainSerializer, PlainValidator, RootModel
 
 from ethereum_test_base_types import HexNumber
+from ethereum_test_fixtures import FIXTURE_FORMATS, FixtureFormat
 
 from .blockchain import EngineFixture as BlockchainEngineFixture
 from .blockchain import Fixture as BlockchainFixture
@@ -25,7 +26,11 @@ class TestCaseBase(BaseModel):
     id: str
     fixture_hash: HexNumber | None
     fork: str | None
-    format: str
+    format: Annotated[
+        FixtureFormat,
+        PlainSerializer(lambda f: f.fixture_format_name),
+        PlainValidator(lambda f: FIXTURE_FORMATS[f] if f in FIXTURE_FORMATS else f),
+    ]
     __test__ = False  # stop pytest from collecting this class as a test
 
 
diff --git a/src/pytest_plugins/consume/direct/test_via_direct.py b/src/pytest_plugins/consume/direct/test_via_direct.py
index 7b52cf3461..0bdfe13731 100644
--- a/src/pytest_plugins/consume/direct/test_via_direct.py
+++ b/src/pytest_plugins/consume/direct/test_via_direct.py
@@ -9,7 +9,7 @@
 
 import pytest
 
-from ethereum_test_fixtures import FIXTURE_FORMATS, BlockchainFixture, StateFixture
+from ethereum_test_fixtures import BlockchainFixture, StateFixture
 from ethereum_test_fixtures.consume import TestCaseIndexFile, TestCaseStream
 from evm_transition_tool import TransitionTool
 
@@ -30,7 +30,7 @@ def test_blocktest(  # noqa: D103
     if evm_run_single_test:
         fixture_name = re.escape(test_case.id)
     evm.verify_fixture(
-        FIXTURE_FORMATS[test_case.format],
+        test_case.format,
         fixture_path,
         fixture_name=fixture_name,
         debug_output_path=test_dump_dir,
@@ -51,7 +51,7 @@ def run_statetest(
     # TODO: Can we group the tests appropriately so that this works more efficiently with xdist?
     if fixture_path not in statetest_results:
         json_result = evm.verify_fixture(
-            FIXTURE_FORMATS[test_case.format],
+            test_case.format,
             fixture_path,
             fixture_name=None,
             debug_output_path=test_dump_dir,