Joseph Doherty
22 lines
638 B
Python
22 lines
638 B
Python
import pytest
|
|
from chat.llm.mock import MockLLMClient
|
|
from chat.llm.client import Message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mock_returns_canned_response():
|
|
client = MockLLMClient(canned=["Hello, world."])
|
|
msgs = [Message(role="user", content="hi")]
|
|
out = await client.generate(msgs, model="any")
|
|
assert out == "Hello, world."
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mock_streams_tokens():
|
|
client = MockLLMClient(canned=["abcd"])
|
|
msgs = [Message(role="user", content="hi")]
|
|
chunks = []
|
|
async for chunk in client.stream(msgs, model="any"):
|
|
chunks.append(chunk)
|
|
assert "".join(chunks) == "abcd"
|