presenton--presenton/servers/fastapi/tests/test_image_generation_openai_compatible.py

import pytest
import os
from unittest.mock import Mock, patch, AsyncMock, MagicMock
from services.image_generation_service import ImageGenerationService


class TestImageGenerationOpenAICompatible:
    @pytest.fixture
    def anyio_backend(self):
        return "asyncio"

    @pytest.fixture
    def mock_images_directory(self, tmp_path):
        images_dir = tmp_path / "images"
        images_dir.mkdir()
        return str(images_dir)

    def test_get_image_gen_func_openai_compatible_selected(self, mock_images_directory):
        with patch(
            "services.image_generation_service.is_image_generation_disabled",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_openai_compatible_selected",
            return_value=True,
        ), patch(
            "services.image_generation_service.is_pixabay_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_pixels_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_gemini_flash_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_nanobanana_pro_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_dalle3_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_gpt_image_1_5_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_comfyui_selected",
            return_value=False,
        ), patch(
            "services.image_generation_service.is_open_webui_selected",
            return_value=False,
        ):
            service = ImageGenerationService(mock_images_directory)

        assert service.image_gen_func == service.generate_image_openai_compatible

    @pytest.mark.anyio
    async def test_generate_image_openai_compatible_success(
        self, mock_images_directory
    ):
        service = ImageGenerationService(mock_images_directory)

        # Mock environment variables
        with patch(
            "services.image_generation_service.get_openai_compat_image_base_url_env",
            return_value="https://api.example.com/v1",
        ):
            with patch(
                "services.image_generation_service.get_openai_compat_image_api_key_env",
                return_value="sk-test-key",
            ):
                with patch(
                    "services.image_generation_service.get_openai_compat_image_model_env",
                    return_value="custom-model",
                ):
                    # Mock AsyncOpenAI client
                    with patch(
                        "services.image_generation_service.AsyncOpenAI"
                    ) as MockClient:
                        mock_client_instance = MockClient.return_value

                        # Mock response with b64_json
                        mock_response = Mock()
                        mock_data = Mock()
                        # Create a valid base64 image (1x1 pixel transparent png)
                        b64_image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
                        mock_data.b64_json = b64_image
                        mock_data.url = None
                        mock_response.data = [mock_data]

                        mock_client_instance.images.generate = AsyncMock(
                            return_value=mock_response
                        )

                        image_path = await service.generate_image_openai_compatible(
                            "test prompt", mock_images_directory
                        )

                        # Verify client initialization
                        MockClient.assert_called_with(
                            base_url="https://api.example.com/v1", api_key="sk-test-key"
                        )

                        # Verify generate call (no response_format sent)
                        mock_client_instance.images.generate.assert_called_with(
                            model="custom-model",
                            prompt="test prompt",
                            n=1,
                            size="1024x1024",
                        )

                        # Verify file creation
                        assert os.path.exists(image_path)
                        assert image_path.startswith(mock_images_directory)

    @pytest.mark.anyio
    async def test_generate_image_openai_compatible_missing_config(
        self, mock_images_directory
    ):
        service = ImageGenerationService(mock_images_directory)

        with patch(
            "services.image_generation_service.get_openai_compat_image_base_url_env",
            return_value=None,
        ):
            with pytest.raises(
                ValueError,
                match="OPENAI_COMPAT_IMAGE_BASE_URL, OPENAI_COMPAT_IMAGE_API_KEY and OPENAI_COMPAT_IMAGE_MODEL must be set",
            ):
                await service.generate_image_openai_compatible(
                    "test prompt", mock_images_directory
                )

    @pytest.mark.anyio
    async def test_generate_image_openai_compatible_missing_model(
        self, mock_images_directory
    ):
        service = ImageGenerationService(mock_images_directory)

        with patch(
            "services.image_generation_service.get_openai_compat_image_base_url_env",
            return_value="https://api.example.com/v1",
        ):
            with patch(
                "services.image_generation_service.get_openai_compat_image_api_key_env",
                return_value="sk-test-key",
            ):
                with patch(
                    "services.image_generation_service.get_openai_compat_image_model_env",
                    return_value=None,
                ):
                    with pytest.raises(
                        ValueError,
                        match="OPENAI_COMPAT_IMAGE_BASE_URL, OPENAI_COMPAT_IMAGE_API_KEY and OPENAI_COMPAT_IMAGE_MODEL must be set",
                    ):
                        await service.generate_image_openai_compatible(
                            "test prompt", mock_images_directory
                        )

    @pytest.mark.anyio
    async def test_generate_image_openai_compatible_url_response(
        self, mock_images_directory
    ):
        """Providers that return a URL instead of b64_json should be handled correctly."""
        service = ImageGenerationService(mock_images_directory)

        with patch(
            "services.image_generation_service.get_openai_compat_image_base_url_env",
            return_value="https://api.example.com/v1",
        ):
            with patch(
                "services.image_generation_service.get_openai_compat_image_api_key_env",
                return_value="sk-test-key",
            ):
                with patch(
                    "services.image_generation_service.get_openai_compat_image_model_env",
                    return_value="custom-model",
                ):
                    with patch(
                        "services.image_generation_service.AsyncOpenAI"
                    ) as MockClient:
                        mock_client_instance = MockClient.return_value
                        mock_data = Mock()
                        mock_data.b64_json = None
                        mock_data.url = "https://api.example.com/images/result.png"
                        mock_response = Mock()
                        mock_response.data = [mock_data]
                        mock_client_instance.images.generate = AsyncMock(
                            return_value=mock_response
                        )

                        fake_image_bytes = b"\x89PNG\r\n\x1a\n"
                        mock_dl_resp = AsyncMock()
                        mock_dl_resp.status = 200
                        mock_dl_resp.read = AsyncMock(return_value=fake_image_bytes)
                        mock_session = AsyncMock()
                        mock_session.get = AsyncMock(return_value=mock_dl_resp)
                        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
                        mock_session.__aexit__ = AsyncMock(return_value=False)

                        with patch(
                            "services.image_generation_service.aiohttp.ClientSession",
                            return_value=mock_session,
                        ):
                            image_path = await service.generate_image_openai_compatible(
                                "test prompt", mock_images_directory
                            )

                        call_kwargs = mock_session.get.call_args.kwargs
                        assert call_kwargs["headers"] == {
                            "Authorization": "Bearer sk-test-key"
                        }
                        assert os.path.exists(image_path)
                        assert image_path.startswith(mock_images_directory)

    @pytest.mark.anyio
    async def test_generate_image_openai_compatible_relative_url_response(
        self, mock_images_directory
    ):
        """Relative URLs in responses are resolved against the configured base URL origin."""
        service = ImageGenerationService(mock_images_directory)

        with patch(
            "services.image_generation_service.get_openai_compat_image_base_url_env",
            return_value="https://api.example.com/v1",
        ):
            with patch(
                "services.image_generation_service.get_openai_compat_image_api_key_env",
                return_value="sk-test-key",
            ):
                with patch(
                    "services.image_generation_service.get_openai_compat_image_model_env",
                    return_value="custom-model",
                ):
                    with patch(
                        "services.image_generation_service.AsyncOpenAI"
                    ) as MockClient:
                        mock_client_instance = MockClient.return_value
                        mock_data = Mock()
                        mock_data.b64_json = None
                        mock_data.url = "/images/result.png"
                        mock_response = Mock()
                        mock_response.data = [mock_data]
                        mock_client_instance.images.generate = AsyncMock(
                            return_value=mock_response
                        )

                        fake_image_bytes = b"\x89PNG\r\n\x1a\n"
                        mock_dl_resp = AsyncMock()
                        mock_dl_resp.status = 200
                        mock_dl_resp.read = AsyncMock(return_value=fake_image_bytes)
                        mock_session = AsyncMock()
                        mock_session.get = AsyncMock(return_value=mock_dl_resp)
                        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
                        mock_session.__aexit__ = AsyncMock(return_value=False)

                        with patch(
                            "services.image_generation_service.aiohttp.ClientSession",
                            return_value=mock_session,
                        ) as MockSession:
                            image_path = await service.generate_image_openai_compatible(
                                "test prompt", mock_images_directory
                            )

                        # Relative URL must be resolved against the origin
                        call_args = mock_session.get.call_args
                        called_url = call_args[0][0]
                        assert called_url == "https://api.example.com/images/result.png"
                        assert call_args.kwargs["headers"] == {
                            "Authorization": "Bearer sk-test-key"
                        }
                        assert os.path.exists(image_path)

    @pytest.mark.anyio
    async def test_generate_image_openai_compatible_cross_origin_url_skips_auth(
        self, mock_images_directory
    ):
        """Cross-origin asset URLs should not receive the API bearer token."""
        service = ImageGenerationService(mock_images_directory)

        with patch(
            "services.image_generation_service.get_openai_compat_image_base_url_env",
            return_value="https://api.example.com/v1",
        ):
            with patch(
                "services.image_generation_service.get_openai_compat_image_api_key_env",
                return_value="sk-test-key",
            ):
                with patch(
                    "services.image_generation_service.get_openai_compat_image_model_env",
                    return_value="custom-model",
                ):
                    with patch(
                        "services.image_generation_service.AsyncOpenAI"
                    ) as MockClient:
                        mock_client_instance = MockClient.return_value
                        mock_data = Mock()
                        mock_data.b64_json = None
                        mock_data.url = "https://cdn.example.net/images/result.png"
                        mock_response = Mock()
                        mock_response.data = [mock_data]
                        mock_client_instance.images.generate = AsyncMock(
                            return_value=mock_response
                        )

                        fake_image_bytes = b"\x89PNG\r\n\x1a\n"
                        mock_dl_resp = AsyncMock()
                        mock_dl_resp.status = 200
                        mock_dl_resp.read = AsyncMock(return_value=fake_image_bytes)
                        mock_session = AsyncMock()
                        mock_session.get = AsyncMock(return_value=mock_dl_resp)
                        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
                        mock_session.__aexit__ = AsyncMock(return_value=False)

                        with patch(
                            "services.image_generation_service.aiohttp.ClientSession",
                            return_value=mock_session,
                        ):
                            image_path = await service.generate_image_openai_compatible(
                                "test prompt", mock_images_directory
                            )

                        call_kwargs = mock_session.get.call_args.kwargs
                        assert call_kwargs["headers"] == {}
                        assert os.path.exists(image_path)