77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
|
"""Tests for `app.core.llm.MockLlmClient`."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from app.core.llm import LLMResponseError, MockLlmClient
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_mock_client_returns_replay_response():
|
||
|
|
client = MockLlmClient({
|
||
|
|
"test_stage": [
|
||
|
|
{"message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}
|
||
|
|
]
|
||
|
|
})
|
||
|
|
resp = await client.complete(stage="test_stage", messages=[])
|
||
|
|
assert resp["message"]["content"] == "Hello"
|
||
|
|
assert resp["finish_reason"] == "stop"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_mock_client_advances_on_each_call():
|
||
|
|
client = MockLlmClient({
|
||
|
|
"test_stage": [
|
||
|
|
{"message": {"role": "assistant", "content": "First"}, "finish_reason": "stop"},
|
||
|
|
{"message": {"role": "assistant", "content": "Second"}, "finish_reason": "stop"},
|
||
|
|
]
|
||
|
|
})
|
||
|
|
r1 = await client.complete(stage="test_stage", messages=[])
|
||
|
|
r2 = await client.complete(stage="test_stage", messages=[])
|
||
|
|
assert r1["message"]["content"] == "First"
|
||
|
|
assert r2["message"]["content"] == "Second"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_mock_client_raises_on_exhausted_replay():
|
||
|
|
client = MockLlmClient({
|
||
|
|
"test_stage": [
|
||
|
|
{"message": {"role": "assistant", "content": "Only"}, "finish_reason": "stop"}
|
||
|
|
]
|
||
|
|
})
|
||
|
|
await client.complete(stage="test_stage", messages=[])
|
||
|
|
with pytest.raises(LLMResponseError):
|
||
|
|
await client.complete(stage="test_stage", messages=[])
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_mock_client_records_calls():
|
||
|
|
client = MockLlmClient({
|
||
|
|
"test_stage": [
|
||
|
|
{"message": {"role": "assistant", "content": "X"}, "finish_reason": "stop"}
|
||
|
|
]
|
||
|
|
})
|
||
|
|
await client.complete(stage="test_stage", messages=[{"role": "user", "content": "hi"}],
|
||
|
|
tools=[{"name": "calc"}])
|
||
|
|
assert len(client.recorded_calls) == 1
|
||
|
|
assert client.recorded_calls[0]["stage"] == "test_stage"
|
||
|
|
assert client.recorded_calls[0]["messages"][0]["content"] == "hi"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_mock_client_stream_yields_chunks():
|
||
|
|
client = MockLlmClient({
|
||
|
|
"test_stage": [
|
||
|
|
{"message": {"role": "assistant", "content": "Hello world"}, "finish_reason": "stop"}
|
||
|
|
]
|
||
|
|
})
|
||
|
|
chunks = []
|
||
|
|
async for evt in client.stream_complete(stage="test_stage", messages=[]):
|
||
|
|
chunks.append(evt)
|
||
|
|
# Last chunk must have finish_reason
|
||
|
|
assert chunks[-1]["finish_reason"] == "stop"
|
||
|
|
# Earlier chunks have content deltas
|
||
|
|
contents = [c["delta"].get("content", "") for c in chunks if c["delta"].get("content")]
|
||
|
|
assert "".join(contents) == "Hello world"
|