test: fix test

BerriAI · Jan 29, 2025 · 78cf3f3 · 78cf3f3
1 parent c8f5f06
commit 78cf3f3
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py
@@ -727,7 +727,14 @@ def test_stream_chunk_builder_openai_audio_output_usage():
         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
+            response_usage_value_dict = response_usage_value.model_dump(
+                exclude_none=True
+            )
+            if isinstance(v, dict):
+                for key, value in v.items():
+                    assert response_usage_value_dict[key] == value
+            else:
+                assert response_usage_value_dict == v
         else:
             assert response_usage_value == v