rebase

tests/unit/test_llm_mock.py

@@ -0,0 +1,76 @@
+"""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"