"""Tests for ai_processor module.""" import pytest from unittest.mock import patch, MagicMock from factsdb.ai_processor import AIEndpointClient, AIProcessor from factsdb.config import AIEndpointConfig @pytest.fixture def ai_config(): return AIEndpointConfig(url="http://localhost:8000/v1", auth_token="tok") @pytest.fixture def client(ai_config): return AIEndpointClient(ai_config) class TestAIEndpointClient: def test_headers(self, client): h = client._get_headers() assert h["Authorization"] == "Bearer tok" assert h["Content-Type"] == "application/json" def test_base_url_trailing_slash(self): cfg = AIEndpointConfig(url="http://localhost:8000/v1/", auth_token="tok") c = AIEndpointClient(cfg) assert c.base_url == "http://localhost:8000/v1" class TestSendRequest: def test_success(self, client): mock_resp = MagicMock() mock_resp.json.return_value = {"choices": []} mock_resp.raise_for_status.return_value = None with patch("requests.post", return_value=mock_resp) as mp: result = client.send_request({"model": "m"}) mp.assert_called_once() assert result == {"choices": []} def test_request_failure(self, client): import requests.exceptions with patch("requests.post", side_effect=requests.exceptions.ConnectionError("conn")): with pytest.raises(Exception, match="AI endpoint request failed"): client.send_request({}) class TestExtractFacts: def test_single_chunk(self, client): mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": '{"fact": "F1", "key_entities": ["E"], "key_dates": ["2024"]}'}}] } mock_resp.raise_for_status.return_value = None with patch("requests.post", return_value=mock_resp): result = client.extract_facts("short text", "") assert result["fact"] == "F1" assert "E" in result["key_entities"] assert "2024" in result["key_dates"] def test_chunking(self, client): long_text = "x" * 20000 calls = [] mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": '{"fact": "F1", "key_entities": [], "key_dates": []}'}}] } mock_resp.raise_for_status.return_value = None def side_effect(*a, **kw): calls.append(kw.get("json", {})) return mock_resp with patch("requests.post", side_effect=side_effect): result = client.extract_facts(long_text, "") assert len(calls) >= 2 assert result["fact"] def test_json_code_block(self, client): mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": "```json\n{\"fact\": \"FB\", \"key_entities\": [], \"key_dates\": []}\n```"}}] } mock_resp.raise_for_status.return_value = None with patch("requests.post", return_value=mock_resp): result = client.extract_facts("text", "") assert result["fact"] == "FB" def test_generic_code_block(self, client): mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": "```\n{\"fact\": \"FG\", \"key_entities\": [], \"key_dates\": []}\n```"}}] } mock_resp.raise_for_status.return_value = None with patch("requests.post", return_value=mock_resp): result = client.extract_facts("text", "") assert result["fact"] == "FG" def test_chunk_failure_continues(self, client): long_text = "x" * 20000 mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": '{"fact": "OK", "key_entities": [], "key_dates": []}'}}] } mock_resp.raise_for_status.return_value = None call_count = [0] def side_effect(*a, **kw): call_count[0] += 1 if call_count[0] == 1: raise Exception("fail") return mock_resp with patch("requests.post", side_effect=side_effect): result = client.extract_facts(long_text, "") assert "OK" in result["fact"] def test_custom_prompt(self, client): mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": '{"fact": "X", "key_entities": [], "key_dates": []}'}}] } mock_resp.raise_for_status.return_value = None captured = [] def side_effect(*a, **kw): captured.append(kw.get("json", {})) return mock_resp with patch("requests.post", side_effect=side_effect): client.extract_facts("text", "custom prompt text") assert "custom prompt text" in captured[0]["messages"][0]["content"] class TestAIProcessor: def test_extract_facts_from_text(self, ai_config): mock_resp = MagicMock() mock_resp.json.return_value = { "choices": [{"message": {"content": '{"fact": "FA", "key_entities": ["E1"], "key_dates": []}'}}] } mock_resp.raise_for_status.return_value = None with patch("requests.post", return_value=mock_resp): proc = AIProcessor(ai_config) result = proc.extract_facts_from_text("content") assert result["fact"] == "FA" def test_validate_model_supported(self, ai_config): proc = AIProcessor(ai_config) assert proc.validate_model_support("gpt-oss") is True assert proc.validate_model_support("qwen3") is True def test_validate_model_unsupported(self, ai_config): proc = AIProcessor(ai_config) assert proc.validate_model_support("unknown-model") is False