test: adapt streaming model usage tests to next's always-emit behavior

Author: levilentz

tests/lib/adk/test_tracing_module.py

@@ -113,10 +113,11 @@ async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
         mock_service, module = _make_module()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger") as mock_logger, patch.object(
+            _tracing_mod.workflow, "metric_meter", return_value=mock_meter
+        ):
             mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
             result = await module.start_span(trace_id="trace-123", name="test-span")
 
@@ -127,9 +128,7 @@ async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
             description="Temporal tracing span activities dropped after fail-open",
             unit="1",
         )
-        mock_meter.create_counter.return_value.add.assert_called_once_with(
-            1, {"event_type": "start"}
-        )
+        mock_meter.create_counter.return_value.add.assert_called_once_with(1, {"event_type": "start"})
         mock_helpers.execute_activity.assert_called_once()
         mock_service.start_span.assert_not_called()
 
@@ -138,10 +137,11 @@ async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
         span = _make_span()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger") as mock_logger, patch.object(
+            _tracing_mod.workflow, "metric_meter", return_value=mock_meter
+        ):
             mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
             result = await module.end_span(trace_id="trace-123", span=span)
 
@@ -152,18 +152,16 @@ async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
             description="Temporal tracing span activities dropped after fail-open",
             unit="1",
         )
-        mock_meter.create_counter.return_value.add.assert_called_once_with(
-            1, {"event_type": "end"}
-        )
+        mock_meter.create_counter.return_value.add.assert_called_once_with(1, {"event_type": "end"})
         mock_helpers.execute_activity.assert_called_once()
         mock_service.end_span.assert_not_called()
 
     async def test_context_manager_skips_end_when_temporal_start_fails(self):
         mock_service, module = _make_module()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod.workflow, "logger"):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger"):
             mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
             async with module.span(trace_id="trace-123", name="test-span") as span:
                 assert span is None
@@ -175,8 +173,9 @@ async def test_context_manager_skips_end_when_temporal_start_fails(self):
     async def test_start_span_in_workflow_propagates_unexpected_errors(self):
         mock_service, module = _make_module()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers:
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers:
             mock_helpers.execute_activity = AsyncMock(side_effect=RuntimeError("bad response shape"))
             try:
                 await module.start_span(trace_id="trace-123", name="test-span")
@@ -193,11 +192,11 @@ async def test_start_span_in_workflow_propagates_cancellation(self):
         activity_error = _make_activity_error()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), patch.object(
+            _tracing_mod.workflow, "logger"
+        ) as mock_logger, patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
             mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
 
             with pytest.raises(ActivityError):
@@ -214,11 +213,11 @@ async def test_end_span_in_workflow_propagates_cancellation(self):
         activity_error = _make_activity_error()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), patch.object(
+            _tracing_mod.workflow, "logger"
+        ) as mock_logger, patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
             mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
 
             with pytest.raises(ActivityError):

tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py

@@ -183,7 +183,6 @@ async def test_streaming_model_captures_final_response_usage(self, tracing_conte
         )
 
         # Real usage lands on the returned ModelResponse (was zeroed before)
-        assert response.usage.requests == 1
         assert response.usage.input_tokens == 120
         assert response.usage.output_tokens == 80
         assert response.usage.total_tokens == 200
@@ -195,7 +194,7 @@ async def test_streaming_model_captures_final_response_usage(self, tracing_conte
         span = fake_tracer.trace_obj.spans[0]
         assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
 
-    async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_contextvars):
+    async def test_streaming_model_writes_zero_usage_when_api_reports_none(self, tracing_contextvars):
         completed = ResponseCompletedEvent.model_construct(
             type="response.completed",
             response=Response.model_construct(output=[], usage=None),
@@ -219,6 +218,14 @@ async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_c
             tracing=None,
         )
 
+        # No usage from the API: the model reports zeros rather than omitting,
+        # so billing sums 0 instead of missing the span
         assert response.usage.input_tokens == 0
         span = fake_tracer.trace_obj.spans[0]
-        assert "usage" not in _output_dict(span)
+        assert _output_dict(span)["usage"] == {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "total_tokens": 0,
+            "cached_input_tokens": 0,
+            "reasoning_tokens": 0,
+        }