|
1 |
| -import unittest |
2 | 1 | import uuid
|
3 | 2 | from typing import Optional, Union
|
4 | 3 |
|
@@ -335,54 +334,45 @@ def test_ai_message_chunks() -> None:
|
335 | 334 | )
|
336 | 335 |
|
337 | 336 |
|
338 |
| -class TestGetBufferString(unittest.TestCase): |
339 |
| - def setUp(self) -> None: |
340 |
| - self.human_msg = HumanMessage(content="human") |
341 |
| - self.ai_msg = AIMessage(content="ai") |
342 |
| - self.sys_msg = SystemMessage(content="system") |
343 |
| - self.func_msg = FunctionMessage(name="func", content="function") |
344 |
| - self.tool_msg = ToolMessage(tool_call_id="tool_id", content="tool") |
345 |
| - self.chat_msg = ChatMessage(role="Chat", content="chat") |
346 |
| - self.tool_calls_msg = AIMessage(content="tool") |
| 337 | +class TestGetBufferString: |
| 338 | + _HUMAN_MSG = HumanMessage(content="human") |
| 339 | + _AI_MSG = AIMessage(content="ai") |
347 | 340 |
|
348 | 341 | def test_empty_input(self) -> None:
|
349 | 342 | assert get_buffer_string([]) == ""
|
350 | 343 |
|
351 | 344 | def test_valid_single_message(self) -> None:
|
352 |
| - expected_output = f"Human: {self.human_msg.content}" |
353 |
| - assert get_buffer_string([self.human_msg]) == expected_output |
| 345 | + expected_output = "Human: human" |
| 346 | + assert get_buffer_string([self._HUMAN_MSG]) == expected_output |
354 | 347 |
|
355 | 348 | def test_custom_human_prefix(self) -> None:
|
356 |
| - prefix = "H" |
357 |
| - expected_output = f"{prefix}: {self.human_msg.content}" |
358 |
| - assert get_buffer_string([self.human_msg], human_prefix="H") == expected_output |
| 349 | + expected_output = "H: human" |
| 350 | + assert get_buffer_string([self._HUMAN_MSG], human_prefix="H") == expected_output |
359 | 351 |
|
360 | 352 | def test_custom_ai_prefix(self) -> None:
|
361 |
| - prefix = "A" |
362 |
| - expected_output = f"{prefix}: {self.ai_msg.content}" |
363 |
| - assert get_buffer_string([self.ai_msg], ai_prefix="A") == expected_output |
| 353 | + expected_output = "A: ai" |
| 354 | + assert get_buffer_string([self._AI_MSG], ai_prefix="A") == expected_output |
364 | 355 |
|
365 | 356 | def test_multiple_msg(self) -> None:
|
366 | 357 | msgs = [
|
367 |
| - self.human_msg, |
368 |
| - self.ai_msg, |
369 |
| - self.sys_msg, |
370 |
| - self.func_msg, |
371 |
| - self.tool_msg, |
372 |
| - self.chat_msg, |
373 |
| - self.tool_calls_msg, |
| 358 | + self._HUMAN_MSG, |
| 359 | + self._AI_MSG, |
| 360 | + SystemMessage(content="system"), |
| 361 | + FunctionMessage(name="func", content="function"), |
| 362 | + ToolMessage(tool_call_id="tool_id", content="tool"), |
| 363 | + ChatMessage(role="Chat", content="chat"), |
| 364 | + AIMessage(content="tool"), |
374 | 365 | ]
|
375 |
| - expected_output = "\n".join( # noqa: FLY002 |
376 |
| - [ |
377 |
| - "Human: human", |
378 |
| - "AI: ai", |
379 |
| - "System: system", |
380 |
| - "Function: function", |
381 |
| - "Tool: tool", |
382 |
| - "Chat: chat", |
383 |
| - "AI: tool", |
384 |
| - ] |
| 366 | + expected_output = ( |
| 367 | + "Human: human\n" |
| 368 | + "AI: ai\n" |
| 369 | + "System: system\n" |
| 370 | + "Function: function\n" |
| 371 | + "Tool: tool\n" |
| 372 | + "Chat: chat\n" |
| 373 | + "AI: tool" |
385 | 374 | )
|
| 375 | + |
386 | 376 | assert get_buffer_string(msgs) == expected_output
|
387 | 377 |
|
388 | 378 |
|
|
0 commit comments