Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| from unittest.mock import MagicMock | |
| import pytest | |
| from bl_ocr.core.config import Settings | |
| from bl_ocr.core.errors import OcrProviderError | |
| from bl_ocr.ocr.client import MistralOcrClient | |
| from bl_ocr.ocr.prompt import PromptAssets | |
| def make_settings( | |
| url: str = "https://api.mistral.ai/v1/chat/completions", | |
| model: str = "ministral-8b-2512", | |
| api_key: str | None = "fake-key", | |
| ) -> Settings: | |
| return Settings( | |
| service_name="test", | |
| api_version="v1", | |
| mistral_api_key=api_key, | |
| mistral_ocr_url=url, | |
| mistral_model=model, | |
| mistral_strict_schema=True, | |
| request_timeout_seconds=5.0, | |
| max_retries=1, | |
| max_upload_mb=1, | |
| schema_path=Path("schema.json"), | |
| example_path=Path("example.json"), | |
| cors_allow_origins=("*",), | |
| ) | |
| def test_client_requires_api_key(): | |
| settings = make_settings(api_key=None) | |
| assets = PromptAssets.load(settings) | |
| client = MistralOcrClient(settings=settings, assets=assets) | |
| with pytest.raises(OcrProviderError, match="MISTRAL_API_KEY is not configured"): | |
| client.extract(b"dummy-image", "image/jpeg") | |
| def test_client_formats_chat_completions_payload(): | |
| settings = make_settings( | |
| url="https://api.mistral.ai/v1/chat/completions", | |
| model="ministral-8b-2512", | |
| ) | |
| assets = PromptAssets.load(settings) | |
| mock_session = MagicMock() | |
| mock_resp = MagicMock() | |
| mock_resp.status_code = 200 | |
| expected_content = json.dumps({"supplier": {"name": {"value": "PHARMA"}}}) | |
| mock_resp.json.return_value = { | |
| "choices": [{"message": {"content": expected_content}}], | |
| "id": "chat-123", | |
| } | |
| mock_session.post.return_value = mock_resp | |
| client = MistralOcrClient(settings=settings, assets=assets, session=mock_session) | |
| result = client.extract(b"image-bytes", "image/jpeg") | |
| assert result["document_annotation"] == expected_content | |
| assert result["raw_chat_response"]["id"] == "chat-123" | |
| mock_session.post.assert_called_once() | |
| called_url = mock_session.post.call_args[0][0] | |
| called_kwargs = mock_session.post.call_args[1] | |
| assert called_url == "https://api.mistral.ai/v1/chat/completions" | |
| assert called_kwargs["headers"]["Authorization"] == "Bearer fake-key" | |
| payload = called_kwargs["json"] | |
| assert payload["model"] == "ministral-8b-2512" | |
| assert "messages" in payload | |
| assert len(payload["messages"]) == 1 | |
| assert payload["messages"][0]["role"] == "user" | |
| assert payload["response_format"]["type"] == "json_schema" | |
| assert payload["response_format"]["json_schema"]["strict"] is True | |
| def test_client_formats_dedicated_ocr_payload(): | |
| settings = make_settings( | |
| url="https://api.mistral.ai/v1/ocr", | |
| model="mistral-ocr-latest", | |
| ) | |
| assets = PromptAssets.load(settings) | |
| mock_session = MagicMock() | |
| mock_resp = MagicMock() | |
| mock_resp.status_code = 200 | |
| mock_resp.json.return_value = { | |
| "pages": [], | |
| "document_annotation": '{"supplier": {}}', | |
| } | |
| mock_session.post.return_value = mock_resp | |
| client = MistralOcrClient(settings=settings, assets=assets, session=mock_session) | |
| result = client.extract(b"image-bytes", "image/jpeg") | |
| assert result["document_annotation"] == '{"supplier": {}}' | |
| called_kwargs = mock_session.post.call_args[1] | |
| payload = called_kwargs["json"] | |
| assert payload["model"] == "mistral-ocr-latest" | |
| assert "document" in payload | |
| assert payload["document"]["type"] == "image_url" | |
| assert "document_annotation_format" in payload | |
| def test_client_handles_http_errors(): | |
| settings = make_settings() | |
| assets = PromptAssets.load(settings) | |
| mock_session = MagicMock() | |
| mock_resp = MagicMock() | |
| mock_resp.status_code = 429 | |
| mock_resp.text = '{"error": "rate limit"}' | |
| mock_session.post.return_value = mock_resp | |
| client = MistralOcrClient(settings=settings, assets=assets, session=mock_session) | |
| with pytest.raises(OcrProviderError) as exc_info: | |
| client.extract(b"image-bytes", "image/jpeg") | |
| assert exc_info.value.status_code == 429 | |