fastgpt-python-sdk/tests/test_image_inputs.py

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