"""Tests for image-input helpers and the base64/upload handling modes.""" from unittest.mock import AsyncMock, Mock, patch import httpx import pytest from fastgpt_client import images from fastgpt_client.async_client import AsyncChatClient from fastgpt_client.client import ChatClient from fastgpt_client.exceptions import ValidationError class TestImageHelpers: def test_encode_decode_round_trip(self): url = images.encode_image_data_url(b"hello", "image/png") assert url.startswith("data:image/png;base64,") assert images.is_data_url(url) mime, raw = images.decode_data_url(url) assert mime == "image/png" assert raw == b"hello" def test_is_data_url_rejects_plain_url(self): assert not images.is_data_url("https://example.com/a.png") assert not images.is_data_url(None) def test_decode_rejects_non_data_url(self): with pytest.raises(ValueError): images.decode_data_url("https://example.com/a.png") def test_decode_rejects_bad_base64(self): with pytest.raises(ValueError): images.decode_data_url("data:image/png;base64,!!!notbase64!!!") def test_image_url_part_shape(self): assert images.image_url_part("u") == { "type": "image_url", "image_url": {"url": "u"}, } def test_image_part_from_path(self, tmp_path): f = tmp_path / "pic.png" f.write_bytes(b"\x89PNG\r\n") part = images.image_part_from_path(f) assert part["type"] == "image_url" assert part["image_url"]["url"].startswith("data:image/png;base64,") def _data_url(): return images.encode_image_data_url(b"imgbytes", "image/jpeg") class TestSyncImageInputMode: def test_base64_mode_passes_through(self, api_key): client = ChatClient(api_key) url = _data_url() messages = [{"role": "user", "content": [images.image_url_part(url)]}] mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 with patch.object(client, "_send_request", return_value=mock_response) as send, \ patch.object(client, "upload_chat_image") as upload: client.create_chat_completion(messages=messages) upload.assert_not_called() sent = send.call_args[1]["json"]["messages"] assert sent[0]["content"][0]["image_url"]["url"] == url def test_upload_mode_replaces_data_url(self, api_key): client = ChatClient(api_key) url = _data_url() messages = [ { "role": "user", "content": [ {"type": "text", "text": "hi"}, images.image_url_part(url), images.image_url_part("https://keep/me.png"), ], } ] mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 with patch.object(client, "_send_request", return_value=mock_response) as send, \ patch.object( client, "upload_chat_image", return_value={"url": "https://cdn/up.jpg"} ) as upload: client.create_chat_completion( messages=messages, chatId="chat-1", appId="app-1", image_input_mode="upload", ) upload.assert_called_once() sent = send.call_args[1]["json"]["messages"] assert sent[0]["content"][1]["image_url"]["url"] == "https://cdn/up.jpg" # plain URL part untouched assert sent[0]["content"][2]["image_url"]["url"] == "https://keep/me.png" # original messages not mutated assert messages[0]["content"][1]["image_url"]["url"] == url def test_upload_mode_requires_app_and_chat_id(self, api_key): client = ChatClient(api_key) messages = [{"role": "user", "content": [images.image_url_part(_data_url())]}] with pytest.raises(ValidationError): client.create_chat_completion(messages=messages, image_input_mode="upload") def test_invalid_mode_raises(self, api_key): client = ChatClient(api_key) messages = [{"role": "user", "content": "hi"}] with pytest.raises(ValidationError): client.create_chat_completion(messages=messages, image_input_mode="nope") def test_upload_failure_falls_back_to_base64(self, api_key): client = ChatClient(api_key) url = _data_url() messages = [{"role": "user", "content": [images.image_url_part(url)]}] mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 with patch.object(client, "_send_request", return_value=mock_response) as send, \ patch.object(client, "upload_chat_image", side_effect=RuntimeError("boom")): client.create_chat_completion( messages=messages, chatId="c", appId="a", image_input_mode="upload" ) sent = send.call_args[1]["json"]["messages"] assert sent[0]["content"][0]["image_url"]["url"] == url class TestAsyncImageInputMode: @pytest.mark.asyncio async def test_upload_mode_replaces_data_url(self, api_key): client = AsyncChatClient(api_key) url = _data_url() messages = [{"role": "user", "content": [images.image_url_part(url)]}] mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 with patch.object(client, "_send_request", new=AsyncMock(return_value=mock_response)) as send, \ patch.object( client, "upload_chat_image", new=AsyncMock(return_value={"url": "https://cdn/u.jpg"}), ): await client.create_chat_completion( messages=messages, chatId="c", appId="a", image_input_mode="upload" ) sent = send.call_args[1]["json"]["messages"] assert sent[0]["content"][0]["image_url"]["url"] == "https://cdn/u.jpg" await client.close() @pytest.mark.asyncio async def test_upload_mode_requires_app_and_chat_id(self, api_key): client = AsyncChatClient(api_key) messages = [{"role": "user", "content": [images.image_url_part(_data_url())]}] with pytest.raises(ValidationError): await client.create_chat_completion(messages=messages, image_input_mode="upload") await client.close()