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.
131 lines
4.6 KiB
Python
131 lines
4.6 KiB
Python
"""Unit tests for OllamaClient."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from libs.common.retries import RetryableError
|
|
from libs.llm.client import OllamaClient
|
|
from libs.llm.exceptions import LLMError, 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")
|
|
|
|
@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"}
|
|
sync_result = (expected, {"prompt_tokens": 100, "completion_tokens": 50}, 42)
|
|
|
|
with patch.object(client, "_sync_call", return_value=sync_result):
|
|
async with client:
|
|
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)."""
|
|
with patch.object(
|
|
client,
|
|
"_sync_call",
|
|
side_effect=LLMTimeoutError("timeout", source="ollama", context={}),
|
|
):
|
|
async with client:
|
|
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."""
|
|
with patch.object(
|
|
client,
|
|
"_sync_call",
|
|
side_effect=RetryableError(
|
|
"Ollama server error 503",
|
|
source="ollama",
|
|
context={"status": 503, "body": "Service Unavailable"},
|
|
),
|
|
):
|
|
async with client:
|
|
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()
|