diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py index 2c3eb20406..8e7cfcf9ed 100644 --- a/tests/local_testing/test_stream_chunk_builder.py +++ b/tests/local_testing/test_stream_chunk_builder.py @@ -721,15 +721,14 @@ def test_stream_chunk_builder_openai_audio_output_usage(): print(f"response usage: {response.usage}") check_non_streaming_response(response) print(f"response: {response}") - for k, v in usage_obj.model_dump(exclude_none=True).items(): - print(k, v) - response_usage_value = getattr(response.usage, k) # type: ignore - print(f"response_usage_value: {response_usage_value}") - print(f"type: {type(response_usage_value)}") - if isinstance(response_usage_value, BaseModel): - assert response_usage_value.model_dump(exclude_none=True) == v - else: - assert response_usage_value == v + # Convert both usage objects to dictionaries for easier comparison + usage_dict = usage_obj.model_dump(exclude_none=True) + response_usage_dict = response.usage.model_dump(exclude_none=True) + + # Simple dictionary comparison + assert ( + usage_dict == response_usage_dict + ), f"\nExpected: {usage_dict}\nGot: {response_usage_dict}" def test_stream_chunk_builder_empty_initial_chunk():