You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
138 lines
4.8 KiB
Python
138 lines
4.8 KiB
Python
"""Unit tests for OllamaClient."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from libs.common.retries import RetryableError
|
|
from libs.llm.client import OllamaClient
|
|
from libs.llm.exceptions import LLMTimeoutError
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestOllamaClientChat:
|
|
"""Tests for OllamaClient.chat()."""
|
|
|
|
@pytest.fixture
|
|
def client(self) -> OllamaClient:
|
|
return OllamaClient(base_url="http://localhost:11434", model="llama3.2")
|
|
|
|
async def _mock_response(self, payload: dict) -> httpx.Response:
|
|
return httpx.Response(200, json=payload)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_success_returns_parsed_json(self, client: OllamaClient) -> None:
|
|
"""A successful Ollama response is parsed and returned as a dict."""
|
|
expected = {"event_type": "earnings_release", "event_direction": "bullish"}
|
|
mock_resp_payload = {
|
|
"message": {"content": json.dumps(expected)},
|
|
"prompt_eval_count": 100,
|
|
"eval_count": 50,
|
|
}
|
|
|
|
mock_http = AsyncMock()
|
|
mock_http.post = AsyncMock(return_value=httpx.Response(200, json=mock_resp_payload))
|
|
|
|
async with client:
|
|
client._client = mock_http # inject mock
|
|
result, token_usage, elapsed_ms = await client.chat(
|
|
[{"role": "user", "content": "analyze this"}]
|
|
)
|
|
|
|
assert result == expected
|
|
assert token_usage["prompt_tokens"] == 100
|
|
assert token_usage["completion_tokens"] == 50
|
|
assert elapsed_ms >= 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_timeout_raises_llm_timeout_error(self, client: OllamaClient) -> None:
|
|
"""Timeout raises LLMTimeoutError (which is also RetryableError)."""
|
|
mock_http = AsyncMock()
|
|
mock_http.post = AsyncMock(side_effect=httpx.ReadTimeout("timeout"))
|
|
|
|
async with client:
|
|
client._client = mock_http
|
|
with pytest.raises(LLMTimeoutError):
|
|
await client.chat([{"role": "user", "content": "test"}])
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_5xx_raises_retryable_error(self, client: OllamaClient) -> None:
|
|
"""5xx response raises RetryableError."""
|
|
mock_http = AsyncMock()
|
|
mock_http.post = AsyncMock(
|
|
return_value=httpx.Response(503, text="Service Unavailable")
|
|
)
|
|
|
|
async with client:
|
|
client._client = mock_http
|
|
with pytest.raises(RetryableError):
|
|
await client.chat([{"role": "user", "content": "test"}])
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_hit_skips_llm_call(self) -> None:
|
|
"""LLMParser returns cached result without calling Ollama."""
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
from libs.llm.cache import LLMCacheStore
|
|
from libs.llm.parser import LLMParser
|
|
|
|
cached_output = {
|
|
"schema_version": "1.0.0",
|
|
"document_id": "DOC::test",
|
|
"parser_kind": "llm",
|
|
"event_type": "earnings_release",
|
|
"event_direction": "bullish",
|
|
"event_date": "2026-01-01",
|
|
"filing_time_bucket": "post_market",
|
|
"headline": "Test",
|
|
"summary": "Test summary",
|
|
"guidance": {"status": "raised", "scope": "annual", "notes": ""},
|
|
"signals": {
|
|
"demand_strength": "strong",
|
|
"pricing_power": "present",
|
|
"backlog_or_bookings": "present",
|
|
"customer_expansion": "present",
|
|
"margin_quality": "improving",
|
|
},
|
|
"risk_flags": {
|
|
"oneoff_item": False,
|
|
"tax_benefit": False,
|
|
"valuation_gain": False,
|
|
"non_gaap_heavy": False,
|
|
"financing_related": False,
|
|
"legal_or_regulatory_overhang": False,
|
|
},
|
|
"evidence": [],
|
|
"confidence": {
|
|
"overall": 0.85,
|
|
"event_type": 0.9,
|
|
"event_direction": 0.8,
|
|
"guidance": 0.8,
|
|
"risk_flags": 0.9,
|
|
},
|
|
"warnings": [],
|
|
}
|
|
|
|
mock_cache = MagicMock(spec=LLMCacheStore)
|
|
mock_cache.get = AM(return_value=cached_output)
|
|
|
|
mock_client = MagicMock(spec=OllamaClient)
|
|
mock_client.model = "llama3.2"
|
|
|
|
llm_parser = LLMParser(client=mock_client, cache_store=mock_cache)
|
|
mock_session = AsyncMock()
|
|
|
|
result = await llm_parser.parse(
|
|
document_id="DOC::test",
|
|
doc_text="some text",
|
|
doc_meta={"filing_date": "2026-01-01", "form_type": "8-K"},
|
|
rule_hints={},
|
|
session=mock_session,
|
|
)
|
|
|
|
assert result == cached_output
|
|
mock_client.chat.assert_not_called()
|