From b3448f4baa75a335b111ad517e793a110ad449ff Mon Sep 17 00:00:00 2001
From: Lina Tawfik <lina@anthropic.com>
Date: Tue, 17 Jun 2025 16:51:28 -0700
Subject: [PATCH] Fix tests to match updated ResultMessage type
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Remove cost_usd field from test expectations
- Update to use total_cost_usd instead of total_cost
- All 30 tests now pass

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 tests/test_client.py      | 3 +--
 tests/test_integration.py | 8 +++-----
 tests/test_types.py       | 3 +--
 3 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/tests/test_client.py b/tests/test_client.py
index 8ade785..3282ea1 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -88,13 +88,12 @@ class TestQueryFunction:
                     yield {
                         "type": "result",
                         "subtype": "success",
-                        "cost_usd": 0.001,
                         "duration_ms": 1000,
                         "duration_api_ms": 800,
                         "is_error": False,
                         "num_turns": 1,
                         "session_id": "test-session",
-                        "total_cost": 0.001,
+                        "total_cost_usd": 0.001,
                     }
 
                 mock_transport.receive_messages = mock_receive
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 0798569..a185335 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -43,13 +43,12 @@ class TestIntegration:
                     yield {
                         "type": "result",
                         "subtype": "success",
-                        "cost_usd": 0.001,
                         "duration_ms": 1000,
                         "duration_api_ms": 800,
                         "is_error": False,
                         "num_turns": 1,
                         "session_id": "test-session",
-                        "total_cost": 0.001,
+                        "total_cost_usd": 0.001,
                     }
 
                 mock_transport.receive_messages = mock_receive
@@ -71,7 +70,7 @@ class TestIntegration:
 
                 # Check result message
                 assert isinstance(messages[1], ResultMessage)
-                assert messages[1].cost_usd == 0.001
+                assert messages[1].total_cost_usd == 0.001
                 assert messages[1].session_id == "test-session"
 
         anyio.run(_test)
@@ -109,13 +108,12 @@ class TestIntegration:
                     yield {
                         "type": "result",
                         "subtype": "success",
-                        "cost_usd": 0.002,
                         "duration_ms": 1500,
                         "duration_api_ms": 1200,
                         "is_error": False,
                         "num_turns": 1,
                         "session_id": "test-session-2",
-                        "total_cost": 0.002,
+                        "total_cost_usd": 0.002,
                     }
 
                 mock_transport.receive_messages = mock_receive
diff --git a/tests/test_types.py b/tests/test_types.py
index 7336204..6046292 100644
--- a/tests/test_types.py
+++ b/tests/test_types.py
@@ -45,7 +45,6 @@ class TestMessageTypes:
         """Test creating a ResultMessage."""
         msg = ResultMessage(
             subtype="success",
-            cost_usd=0.01,
             duration_ms=1500,
             duration_api_ms=1200,
             is_error=False,
@@ -54,7 +53,7 @@ class TestMessageTypes:
             total_cost_usd=0.01,
         )
         assert msg.subtype == "success"
-        assert msg.cost_usd == 0.01
+        assert msg.total_cost_usd == 0.01
         assert msg.session_id == "session-123"