From 3e11c986766ed7a32ead704e5284fbfeaf889c6b Mon Sep 17 00:00:00 2001
From: kerry <kerry@berri.ai>
Date: Thu, 17 Sep 2026 01:02:41 +0000
Subject: [PATCH] test(e2e): satisfy pyright in cost matrix derivation and
 golden generator

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 tests/e2e/cost_calculation/cost_matrix.py       | 14 +++++++-------
 tests/e2e/cost_calculation/generate_expected.py | 16 +++++++++++++---
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py
index a8f60b79ae7..35344a099be 100644
--- a/tests/e2e/cost_calculation/cost_matrix.py
+++ b/tests/e2e/cost_calculation/cost_matrix.py
@@ -267,23 +267,23 @@ def _frontier() -> tuple[FrontierModel, ...]:
     )
     models: list[FrontierModel] = []  # mutable-ok: accumulated once at import into a tuple
     for map_key in sorted(COST_MAP):
-        entry: Final = COST_MAP[map_key]
-        pair: Final = (entry.litellm_provider, entry.mode)
-        wiring: Final = _PROVIDER_WIRING.get(pair)
+        entry = COST_MAP[map_key]
+        pair = (entry.litellm_provider, entry.mode)
+        wiring = _PROVIDER_WIRING.get(pair)
         if wiring is None:
             raise ValueError(
                 f"cost_map entry {map_key} has no wiring for "
                 f"(litellm_provider={pair[0]}, mode={pair[1]}); add a "
                 f"_ProviderWiring row in cost_matrix.py"
             )
-        siblings: Final = groups[pair]
-        override_key: Final = (
+        siblings = groups[pair]
+        override_key = (
             siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None
         )
-        override_litellm: Final = (
+        override_litellm = (
             _litellm_model_for(override_key, wiring) if override_key is not None else None
         )
-        deployment: Final = _DEPLOYMENTS.get(map_key)
+        deployment = _DEPLOYMENTS.get(map_key)
         models.append(
             FrontierModel(
                 model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}",
diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py
index e243e477839..f5514092c88 100644
--- a/tests/e2e/cost_calculation/generate_expected.py
+++ b/tests/e2e/cost_calculation/generate_expected.py
@@ -19,6 +19,8 @@ from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final
 
+from pydantic import TypeAdapter
+
 from cost_matrix import (
     EXPECTED_PATH,
     FRONTIER_MODELS,
@@ -168,11 +170,19 @@ def main() -> None:
     rewrite: Final = "--rewrite" in sys.argv[1:]
     proposed: Final = _proposed()
     proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()}
-    existing: Final = (
-        json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {}
+    existing: Final[Mapping[str, ExpectedCell]] = (
+        TypeAdapter(dict[str, ExpectedCell]).validate_python(
+            json.loads(EXPECTED_PATH.read_text())
+        )
+        if EXPECTED_PATH.exists()
+        else {}
     )
     merged: Final = {
-        key: (proposed_values[key] if rewrite or key not in existing else existing[key])
+        key: (
+            proposed_values[key]
+            if rewrite or key not in existing
+            else existing[key].model_dump()
+        )
         for key in sorted(proposed_values)
     }
     added: Final = sum(1 for key in proposed_values if key not in existing)