"""Tests for api_image_request module.""" import json from unittest.mock import MagicMock, patch import pytest from PIL import Image from docling.datamodel.base_models import VlmStopReason from docling.utils.api_image_request import api_image_request pytestmark = pytest.mark.cross_platform class TestApiImageRequest: """Test cases for api_image_request function.""" @pytest.fixture def sample_image(self): """Create a simple test image.""" return Image.new("RGB", (100, 100), color="red") @pytest.fixture def mock_response_factory(self): """Factory to create mock API responses.""" def _create_mock_response( content="Test response", finish_reason="stop", total_tokens=100, status_ok=True, message=None, ): mock_resp = MagicMock() mock_resp.ok = status_ok mock_resp.text = json.dumps( { "id": "test-id", "created": 1234567890, "choices": [ { "index": 0, "message": message or {"role": "assistant", "content": content}, "finish_reason": finish_reason, } ], "usage": { "prompt_tokens": 50, "completion_tokens": 50, "total_tokens": total_tokens, }, } ) return mock_resp return _create_mock_response @patch("docling.utils.api_image_request.requests.post") def test_content_filter_finish_reason( self, mock_post, sample_image, mock_response_factory ): """Test that content_filter finish reason returns CONTENT_FILTERED.""" mock_post.return_value = mock_response_factory( content="Filtered content", finish_reason="content_filter" ) result_text, _tokens, stop_reason = api_image_request( image=sample_image, prompt="Test prompt", url="http://test.api/v1/chat/completions", ) assert result_text == "Filtered content" assert stop_reason == VlmStopReason.CONTENT_FILTERED @patch("docling.utils.api_image_request.requests.post") def test_length_finish_reason(self, mock_post, sample_image, mock_response_factory): """Test that length finish reason returns LENGTH.""" mock_post.return_value = mock_response_factory( content="Truncated content", finish_reason="length" ) result_text, _tokens, stop_reason = api_image_request( image=sample_image, prompt="Test prompt", url="http://test.api/v1/chat/completions", ) assert result_text == "Truncated content" assert stop_reason == VlmStopReason.LENGTH @patch("docling.utils.api_image_request.requests.post") def test_stop_finish_reason(self, mock_post, sample_image, mock_response_factory): """Test that stop finish reason returns END_OF_SEQUENCE.""" mock_post.return_value = mock_response_factory( content="Normal completion", finish_reason="stop" ) result_text, _tokens, stop_reason = api_image_request( image=sample_image, prompt="Test prompt", url="http://test.api/v1/chat/completions", ) assert result_text == "Normal completion" assert stop_reason == VlmStopReason.END_OF_SEQUENCE @patch("docling.utils.api_image_request.requests.post") def test_tool_calls_response(self, mock_post, sample_image, mock_response_factory): """Test that tool calling responses are converted into generated text.""" mock_post.return_value = mock_response_factory( message={ "role": "assistant", "tool_calls": [ { "function": { "name": "markdown_no_bbox", "arguments": json.dumps( [ {"text": "Extracted text"}, {"text": "Second block"}, ] ), } } ], } ) result_text, tokens, stop_reason = api_image_request( image=sample_image, prompt="Test prompt", url="http://test.api/v1/chat/completions", ) assert result_text == "Extracted text\nSecond block" assert tokens == 100 assert stop_reason == VlmStopReason.END_OF_SEQUENCE