From 789f0c61bf31a885ca8fd7f0f951e6078b37b78c Mon Sep 17 00:00:00 2001
From: kerry <kerry@berri.ai>
Date: Sat, 19 Sep 2026 23:12:43 +0000
Subject: [PATCH] test(integration): only accept omitted breakdown components
 when the case expects zero, fix xai output cost

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../cost_calculation/cost_tracking_cases.json     |  2 +-
 .../cost_calculation/test_cost_tracking.py        | 15 +++++++--------
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json
index ead2435f2ac..3a52dfee93a 100644
--- a/tests/integration/cost_calculation/cost_tracking_cases.json
+++ b/tests/integration/cost_calculation/cost_tracking_cases.json
@@ -28262,7 +28262,7 @@
       "expected": {
         "spend": 0.0135964,
         "input_cost": 0.002484,
-        "output_cost": 0.0011224,
+        "output_cost": 0.0011124,
         "tool_usage_cost": 0.01,
         "prompt_tokens": 1840,
         "completion_tokens": 412
diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py
index d97bfe642a2..8c5f306dd9d 100644
--- a/tests/integration/cost_calculation/test_cost_tracking.py
+++ b/tests/integration/cost_calculation/test_cost_tracking.py
@@ -148,16 +148,15 @@ def _assert_breakdown(
     ):
         if expected_component is None:
             continue
-        actual_value: Final = actual_component or 0.0
-        assert approx_equal(actual_value, expected_component), (
-            f"{case.name}: {field} {actual_component} != expected {expected_component}"
-        )
+        omitted_component_allowed: Final = expected_component == 0.0
+        assert (actual_component is None and omitted_component_allowed) or (
+            actual_component is not None and approx_equal(actual_component, expected_component)
+        ), f"{case.name}: {field} {actual_component} != expected {expected_component}"
         if expected.cost_header and case.response.content_type == "application/json":
             header: Final = response.headers.get(header_name)
-            header_value: Final = float(header) if header is not None else 0.0
-            assert approx_equal(header_value, expected_component), (
-                f"{case.name}: {header_name} {header} != expected {expected_component}"
-            )
+            assert (header is None and omitted_component_allowed) or (
+                header is not None and approx_equal(float(header), expected_component)
+            ), f"{case.name}: {header_name} {header} != expected {expected_component}"
     if expected.cost_header and case.response.content_type == "application/json" and any(
         component is not None
         for component in (