FactsDB/tests/test_ai_processor.py

"""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