refactor: Use generics

Author: marioevz

packages/testing/src/execution_testing/client_clis/cli_types.py

@@ -3,7 +3,7 @@
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Any, Dict, List, Self
+from typing import Annotated, Any, Dict, Generic, List, Self, TypeVar
 
 from pydantic import Field, PlainSerializer, PlainValidator
 
@@ -287,43 +287,53 @@ class Result(CamelModel):
     opcode_count: OpcodeCount | None = None
 
 
+JSONDict = Dict[str, Any]
+
+TRaw = TypeVar("TRaw", str, JSONDict)
 
 
 @dataclass(kw_only=True)
-class LazyAlloc:
+class LazyAlloc(Generic[TRaw]):
     """
     Allocation that is lazily loaded from a JSON file.
     """
 
-    json_contents: Dict | None = None
-    str_contents: str | None = None
+    raw: TRaw
     alloc: Alloc | None = None
 
-    def str_contents_from_cache(self) -> str:
-        """
-        Return the cached, unparsed, alloc from the last t8n execution.
-        """
-        assert self.str_contents is not None, "No string cache found"
-        return self.str_contents
-
-    def json_contents_from_cache(self) -> Dict:
-        """
-        Return the cached, unparsed, json dict from the last t8n execution.
-        """
-        assert self.json_contents is not None, "No json cache found"
-        return self.json_contents
+    def validate(self) -> Alloc:
+        """Validate the alloc."""
+        raise NotImplementedError("validate method not implemented.")
 
-    def get_alloc(self) -> Alloc:
+    def get(self) -> Alloc:
         """Model validate the allocation and return it."""
-        if self.alloc is not None:
-            return self.alloc
-        elif self.json_contents is not None:
-            self.alloc = Alloc.model_validate(self.json_contents)
-            return self.alloc
-        elif self.str_contents is not None:
-            self.alloc = Alloc.model_validate_json(self.str_contents)
-            return self.alloc
-        raise ValueError("No alloc found")
+        if self.alloc is None:
+            self.alloc = self.validate()
+        return self.alloc
+
+
+class LazyAllocJson(LazyAlloc[JSONDict]):
+    """
+    Lazy allocation backed by a JSON dict cache.
+
+    Uses Alloc.model_validate on the dict.
+    """
+
+    def validate(self) -> Alloc:
+        """Validate the alloc."""
+        return Alloc.model_validate(self.raw)
+
+
+class LazyAllocStr(LazyAlloc[str]):
+    """
+    Lazy allocation backed by a JSON dict cache.
+
+    Uses Alloc.model_validate on the dict.
+    """
+
+    def validate(self) -> Alloc:
+        """Validate the alloc."""
+        return Alloc.model_validate_json(self.raw)
 
 
 @dataclass
@@ -342,10 +352,12 @@ def to_files(
         Prepare the input in a directory path in the file system for
         consumption by the t8n tool.
         """
-        if isinstance(self.alloc, LazyAlloc):
-            alloc_contents = self.alloc.str_contents_from_cache()
-        else:
+        if isinstance(self.alloc, Alloc):
             alloc_contents = self.alloc.model_dump_json(**model_dump_config)
+        elif isinstance(self.alloc, LazyAllocStr):
+            alloc_contents = self.alloc.raw
+        else:
+            raise Exception(f"Invalid alloc type: {type(self.alloc)}")
 
         env_contents = self.env.model_dump_json(**model_dump_config)
         txs_contents = (
@@ -375,10 +387,12 @@ def to_files(
 
     def model_dump_json(self, **model_dump_config: Any) -> str:
         """Dump the model in string JSON format."""
-        if isinstance(self.alloc, LazyAlloc):
-            alloc_contents = self.alloc.str_contents_from_cache()
-        else:
+        if isinstance(self.alloc, Alloc):
             alloc_contents = self.alloc.model_dump_json(**model_dump_config)
+        elif isinstance(self.alloc, LazyAllocStr):
+            alloc_contents = self.alloc.raw
+        else:
+            raise Exception(f"Invalid alloc type: {type(self.alloc)}")
 
         env_contents = self.env.model_dump_json(**model_dump_config)
         txs_contents = (
@@ -405,12 +419,14 @@ def model_dump_json(self, **model_dump_config: Any) -> str:
     def model_dump(self, mode: str, **model_dump_config: Any) -> Any:
         """Return the validated model."""
         assert mode == "json", f"Mode {mode} not supported."
-        if isinstance(self.alloc, LazyAlloc):
-            alloc_contents = self.alloc.json_contents_from_cache()
-        else:
+        if isinstance(self.alloc, Alloc):
             alloc_contents = self.alloc.model_dump(
                 mode=mode, **model_dump_config
             )
+        elif isinstance(self.alloc, LazyAllocJson):
+            alloc_contents = self.alloc.raw
+        else:
+            raise Exception(f"Invalid alloc type: {type(self.alloc)}")
 
         env_contents = self.env.model_dump(mode=mode, **model_dump_config)
         txs_contents = [
@@ -450,7 +466,7 @@ def model_validate_files(
         result = Result.model_validate_json(
             json_data=result_data, context=context
         )
-        alloc = LazyAlloc(str_contents=alloc_data)
+        alloc = LazyAllocStr(raw=alloc_data)
         output = cls(result=result, alloc=alloc)
         return output
 
@@ -465,7 +481,7 @@ def model_validate(
         result = Result.model_validate(
             obj=response_json["result"], context=context
         )
-        alloc = LazyAlloc(json_contents=response_json["alloc"])
+        alloc = LazyAllocJson(raw=response_json["alloc"])
         output = cls(result=result, alloc=alloc)
         return output
 
@@ -482,7 +498,7 @@ def model_validate_json(
         result = Result.model_validate(
             obj=parsed_json["result"], context=context
         )
-        alloc = LazyAlloc(str_contents=json.dumps(parsed_json["alloc"]))
+        alloc = LazyAllocStr(raw=json.dumps(parsed_json["alloc"]))
         output = cls(result=result, alloc=alloc)
         return output
 

packages/testing/src/execution_testing/client_clis/clis/besu.py

@@ -191,7 +191,7 @@ def evaluate(
             dump_files_to_directory(
                 debug_output_path,
                 {
-                    "output/alloc.json": output.alloc.json_contents_from_cache(),
+                    "output/alloc.json": output.alloc.raw,
                     "output/result.json": output.result.model_dump(
                         mode="json", **model_dump_config
                     ),

packages/testing/src/execution_testing/client_clis/tests/test_execution_specs.py

@@ -182,7 +182,7 @@ def test_evm_t8n(
                 blob_schedule=Berlin.blob_schedule(),
             ),
         )
-        assert to_json(t8n_output.alloc.get_alloc()) == expected.get("alloc")
+        assert to_json(t8n_output.alloc.get()) == expected.get("alloc")
         if isinstance(default_t8n, ExecutionSpecsTransitionTool):
             # The expected output was generated with geth, instead of deleting
             # any info from this expected output, the fields not returned by

packages/testing/src/execution_testing/client_clis/transition_tool.py

@@ -526,7 +526,7 @@ def _evaluate_server(
                 dump_files_to_directory(
                     debug_output_path,
                     {
-                        "input/alloc.json": request_data.input.alloc.json_contents_from_cache()
+                        "input/alloc.json": request_data.input.alloc.raw
                         if isinstance(request_data.input.alloc, LazyAlloc)
                         else request_data.input.alloc.model_dump(
                             mode="json", **model_dump_config
@@ -571,7 +571,7 @@ def _evaluate_server(
                 dump_files_to_directory(
                     debug_output_path,
                     {
-                        "output/alloc.json": output.alloc.json_contents_from_cache(),
+                        "output/alloc.json": output.alloc.raw,
                         "output/result.json": output.result,
                         "output/txs.rlp": str(output.body),
                         "response_info.txt": response_info,
@@ -626,7 +626,7 @@ def _evaluate_stream(
                 dump_files_to_directory(
                     debug_output_path,
                     {
-                        "output/alloc.json": output.alloc.str_contents_from_cache(),
+                        "output/alloc.json": output.alloc.raw,
                         "output/result.json": output.result,
                         "output/txs.rlp": str(output.body),
                     },

packages/testing/src/execution_testing/specs/blockchain.py

@@ -741,7 +741,7 @@ def generate_block_data(
             print_traces(t8n.get_traces())
             pprint(transition_tool_output.result)
             pprint(previous_alloc)
-            pprint(transition_tool_output.alloc.get_alloc())
+            pprint(transition_tool_output.alloc.get())
             raise e
 
         if len(rejected_txs) > 0 and block.exception is None:
@@ -814,13 +814,13 @@ def make_fixture(
             if block.expected_post_state:
                 self.verify_post_state(
                     t8n,
-                    t8n_state=alloc.get_alloc()
+                    t8n_state=alloc.get()
                     if isinstance(alloc, LazyAlloc)
                     else alloc,
                     expected_state=block.expected_post_state,
                 )
         self.check_exception_test(exception=invalid_blocks > 0)
-        alloc = alloc.get_alloc() if isinstance(alloc, LazyAlloc) else alloc
+        alloc = alloc.get() if isinstance(alloc, LazyAlloc) else alloc
         self.verify_post_state(t8n, t8n_state=alloc)
         info = {}
         if self._opcode_count is not None:
@@ -891,7 +891,7 @@ def make_hive_fixture(
             if block.expected_post_state:
                 self.verify_post_state(
                     t8n,
-                    t8n_state=alloc.get_alloc()
+                    t8n_state=alloc.get()
                     if isinstance(alloc, LazyAlloc)
                     else alloc,
                     expected_state=block.expected_post_state,
@@ -906,7 +906,7 @@ def make_hive_fixture(
             " The framework should never try to execute this test case."
         )
 
-        alloc = alloc.get_alloc() if isinstance(alloc, LazyAlloc) else alloc
+        alloc = alloc.get() if isinstance(alloc, LazyAlloc) else alloc
         self.verify_post_state(t8n, t8n_state=alloc)
 
         # Create base fixture data, common to all fixture formats

packages/testing/src/execution_testing/specs/state.py

@@ -167,7 +167,7 @@ def verify_modified_gas_limit(
                 f"Traces are not equivalent (gas_limit={current_gas_limit})"
             )
             return False
-        modified_tool_alloc = modified_tool_output.alloc.get_alloc()
+        modified_tool_alloc = modified_tool_output.alloc.get()
         try:
             self.post.verify_post_alloc(modified_tool_alloc)
         except Exception as e:
@@ -363,7 +363,7 @@ def make_state_test_fixture(
             debug_output_path=self.get_next_transition_tool_output_path(),
             slow_request=self.is_tx_gas_heavy_test(),
         )
-        output_alloc = transition_tool_output.alloc.get_alloc()
+        output_alloc = transition_tool_output.alloc.get()
 
         try:
             self.post.verify_post_alloc(output_alloc)
@@ -391,7 +391,7 @@ def make_state_test_fixture(
                 self._operation_mode == OpMode.OPTIMIZE_GAS_POST_PROCESSING
             )
             base_tool_output = transition_tool_output
-            base_tool_alloc = base_tool_output.alloc.get_alloc()
+            base_tool_alloc = base_tool_output.alloc.get()
             base_tool_result = base_tool_output.result
 
             assert base_tool_result.traces is not None, "Traces not found."