code/ai_finance_agent/backend/tests/test_ai.py

"""外部协议使用 HTTP MockTransport 验证；这些用例不代表已验证真实模型识别精度。"""
import time
import httpx
import pytest
from sqlalchemy import select
from app.domain.models import CacheEntry
from app.infrastructure.ai import ai, AIClient, AIUnavailable, CircuitBreaker, parse_structured
from app.infrastructure.database import Session
from app.shared.config import settings
from app.shared.schemas import RiskAssessment

def test_structured_fences_and_invalid_values():
    raw = '模型说明\n```json\n{"risk_level":2,"risk_score":0.6,"risk_factors":["测试"],"reasoning":"明确测试依据"}\n```'
    assert parse_structured(raw, RiskAssessment).risk_level == 2
    with pytest.raises(AIUnavailable):
        parse_structured('{"risk_level":0,"risk_score":99}', RiskAssessment)

def test_http_schema_key_and_cache(client, monkeypatch):
    monkeypatch.setattr(settings, 'llm_base_url', 'http://contract.test/v1')
    monkeypatch.setattr(settings, 'llm_model', 'contract-risk')
    monkeypatch.setattr(settings, 'ai_api_key', 'contract-key')
    calls = []
    def respond(request):
        calls.append(request)
        assert request.headers['Authorization'] == 'Bearer contract-key'
        return httpx.Response(200, json={'choices': [{'message': {'content': '{"risk_level":1,"risk_score":0.1,"risk_factors":[],"reasoning":"测试依据"}'}}], 'usage': {'total_tokens': 12}})
    async def exercise():
        model = AIClient()
        await model.http.aclose()
        model.http = httpx.AsyncClient(transport=httpx.MockTransport(respond))
        first, usage = await model.structured('测试请求', RiskAssessment)
        second, hit = await model.structured('测试请求', RiskAssessment)
        assert first == second and usage['total_tokens'] == 12 and hit['cache_hit'] is True
        async with Session() as session:
            assert await session.scalar(select(CacheEntry)) is not None
        await model.close()
    client.portal.call(exercise)
    assert len(calls) == 1

def test_real_models_endpoint_uses_configured_key(client, headers, monkeypatch):
    monkeypatch.setattr(settings, 'vlm_base_url', 'http://contract.test/v1')
    monkeypatch.setattr(settings, 'vlm_model', 'test-vision')
    monkeypatch.setattr(settings, 'ai_api_key', 'contract-key')
    def respond(request):
        assert request.headers['Authorization'] == 'Bearer contract-key'
        return httpx.Response(200, json={'data': [{'id': 'test-vision'}]})
    old = ai.http
    monkeypatch.setattr(ai, 'http', httpx.AsyncClient(transport=httpx.MockTransport(respond)))
    try:
        response = client.post('/api/v1/agents/extract/test', headers=headers['carol'])
        assert response.status_code == 200 and response.json()['data']['status'] == 'success'
    finally:
        client.portal.call(ai.http.aclose)
        monkeypatch.setattr(ai, 'http', old)

def test_embeddings_reject_fake_shape_and_fallback(client, monkeypatch):
    monkeypatch.setattr(settings, 'embedding_base_url', 'http://broken-embed.test')
    monkeypatch.setattr(settings, 'llm_base_url', 'http://fallback.test/v1')
    requests = []
    def respond(request):
        requests.append(request)
        if request.url.host == 'broken-embed.test':
            return httpx.Response(503)
        return httpx.Response(200, json={'data': [{'index': 0, 'embedding': [1.0] + [0.0] * 1023}]})
    async def exercise():
        model = AIClient(); await model.http.aclose()
        model.http = httpx.AsyncClient(transport=httpx.MockTransport(respond))
        assert len((await model.embeddings(['真实协议测试']))[0]) == 1024
        await model.close()
        model.http = httpx.AsyncClient(transport=httpx.MockTransport(lambda req: httpx.Response(200, json={'vectors': [[0.1, 0.2]]})))
        with pytest.raises(AIUnavailable):
            await model.embeddings(['维度非法测试'])
        await model.close()
    client.portal.call(exercise)
    assert len(requests) == 2

def test_circuit_half_open_single_probe(client, monkeypatch):
    monkeypatch.setattr(settings, 'ai_failure_threshold', 2)
    monkeypatch.setattr(settings, 'ai_recovery_seconds', 1)
    async def exercise():
        breaker = CircuitBreaker()
        await breaker.enter(); await breaker.finish(False)
        await breaker.enter(); await breaker.finish(False)
        assert breaker.status() == 'open'
        with pytest.raises(AIUnavailable):
            await breaker.enter()
        breaker.opened_at = time.monotonic() - 2
        await breaker.enter()
        assert breaker.status() == 'half_open'
        with pytest.raises(AIUnavailable):
            await breaker.enter()
        await breaker.finish(True)
        assert breaker.status() == 'closed'
    client.portal.call(exercise)