|
1 | 1 | """Tests for verifying that testing utility code works as expected."""
|
2 | 2 |
|
| 3 | +import time |
3 | 4 | from itertools import cycle
|
4 | 5 | from typing import Any, Optional, Union
|
5 | 6 | from uuid import UUID
|
|
9 | 10 | from langchain_core.callbacks.base import AsyncCallbackHandler
|
10 | 11 | from langchain_core.language_models import (
|
11 | 12 | FakeListChatModel,
|
| 13 | + FakeMessagesListChatModel, |
12 | 14 | GenericFakeChatModel,
|
13 | 15 | ParrotFakeChatModel,
|
14 | 16 | )
|
15 |
| -from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage |
| 17 | +from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage |
16 | 18 | from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
|
17 | 19 | from tests.unit_tests.stubs import (
|
18 | 20 | _any_id_ai_message,
|
@@ -230,3 +232,18 @@ def test_fake_list_chat_model_batch() -> None:
|
230 | 232 | fake = FakeListChatModel(responses=["a", "b", "c"])
|
231 | 233 | resp = fake.batch(["1", "2", "3"])
|
232 | 234 | assert resp == expected
|
| 235 | + |
| 236 | + |
| 237 | +def test_fake_messages_list_chat_model_sleep_delay() -> None: |
| 238 | + sleep_time = 0.1 |
| 239 | + model = FakeMessagesListChatModel( |
| 240 | + responses=[AIMessage(content="A"), AIMessage(content="B")], |
| 241 | + sleep=sleep_time, |
| 242 | + ) |
| 243 | + messages = [HumanMessage(content="C")] |
| 244 | + |
| 245 | + start = time.time() |
| 246 | + model.invoke(messages) |
| 247 | + elapsed = time.time() - start |
| 248 | + |
| 249 | + assert elapsed >= sleep_time |
0 commit comments