Add test for AssistantMessage.usage

jackmpcollins · May 16, 2024 · 1877ecf · 1877ecf
1 parent 5ca0246
commit 1877ecf
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/chat_model/test_message.py b/tests/chat_model/test_message.py
@@ -5,6 +5,7 @@
     AssistantMessage,
     FunctionResultMessage,
     Placeholder,
+    Usage,
     UserMessage,
 )
 from magentic.function_call import FunctionCall
@@ -29,6 +30,13 @@ def test_user_message_format():
     assert user_message_formatted == UserMessage("Hello world")
 
 
+def test_assistant_message_usage():
+    assistant_message = AssistantMessage("Hello")
+    assert assistant_message.usage is None
+    assistant_message._usage_pointer = [Usage(input_tokens=1, output_tokens=2)]
+    assert assistant_message.usage == Usage(input_tokens=1, output_tokens=2)
+
+
 def test_assistant_message_format_str():
     assistant_message = AssistantMessage("Hello {x}")
     assistant_message_formatted = assistant_message.format(x="world")