backend/tests/test_agent_tools.py - main - githubmirror/screenshot-to-code

backend/tests/test_agent_tools.py
314 строк11 KB
Flightdeck Agent
Batch image editing and background removal tools in parallel
30 июл 2026, 18:17
30 июл 2026, 18:17d026163
from typing import Any, cast

import pytest

from agent.state import AgentFileState
from agent.providers.factory import create_provider_session
from agent.tools import canonical_tool_definitions, summarize_tool_input
from agent.tools.types import ToolCall
from image_generation.replicate import P_IMAGE_EDIT_ASPECT_RATIOS
from llm import Llm


def test_canonical_tool_definitions_include_generate_images_when_enabled() -> None:
    tool_names = [tool.name for tool in canonical_tool_definitions(True)]
    assert "generate_images" in tool_names


def test_canonical_tool_definitions_use_plural_batch_image_tool_names() -> None:
    tools = canonical_tool_definitions(True)
    tool_names = [tool.name for tool in tools]

    assert "remove_backgrounds" in tool_names
    assert "edit_images" in tool_names
    assert "remove_background" not in tool_names
    assert "edit_image" not in tool_names
    remove_backgrounds_tool = next(
        tool for tool in tools if tool.name == "remove_backgrounds"
    )
    assert remove_backgrounds_tool.parameters["required"] == ["image_urls"]
    assert (
        remove_backgrounds_tool.parameters["properties"]["image_urls"]["type"]
        == "array"
    )


def test_canonical_tool_definitions_exclude_generate_images_when_disabled() -> None:
    tool_names = [tool.name for tool in canonical_tool_definitions(False)]
    assert "generate_images" not in tool_names


def test_edit_file_tool_description_matches_runtime_output_shape() -> None:
    edit_tool = next(
        tool for tool in canonical_tool_definitions(True) if tool.name == "edit_file"
    )

    assert "success message" in edit_tool.description
    assert "unified diff" in edit_tool.description


def test_canonical_tool_definitions_include_save_assets() -> None:
    tools = canonical_tool_definitions(True)
    tool_names = [tool.name for tool in tools]
    assert "save_assets" in tool_names
    save_assets_tool = next(tool for tool in tools if tool.name == "save_assets")
    assert save_assets_tool.parameters["required"] == ["asset_ids"]
    assert save_assets_tool.parameters["properties"]["asset_ids"]["type"] == "array"


def test_save_assets_tool_input_summary_uses_asset_ids() -> None:
    summary = summarize_tool_input(
        ToolCall(
            id="call-1",
            name="save_assets",
            arguments={"asset_ids": ["asset_one", "asset_two"]},
        ),
        AgentFileState(),
    )

    assert summary == {
        "count": 2,
        "asset_ids": ["asset_one", "asset_two"],
    }


def test_canonical_tool_definitions_include_extract_assets() -> None:
    tools = canonical_tool_definitions(True, asset_extraction_enabled=True)
    tool_names = [tool.name for tool in tools]
    assert "extract_assets" in tool_names
    extract_assets_tool = next(tool for tool in tools if tool.name == "extract_assets")
    assert extract_assets_tool.parameters["required"] == ["asset_descriptions"]
    asset_descriptions = extract_assets_tool.parameters["properties"][
        "asset_descriptions"
    ]
    assert asset_descriptions["type"] == "array"

    # Good caller-side descriptors are essential for separating lookalikes and
    # mapping assets across multiple screenshots instead of making Gemini guess.
    item_description = asset_descriptions["items"]["description"].lower()
    for required_detail in (
        "distinctive appearance",
        "precise location",
        "nearby ui/context",
        "1-based screenshot number",
        "repeated lookalikes",
        "leftmost vs. rightmost",
    ):
        assert required_detail in item_description

    tool_description = extract_assets_tool.description.lower()
    assert "request order" in tool_description
    assert "permanent, embeddable public_url" in tool_description
    assert "attached crop preview" in tool_description
    assert "absent or unisolatable" in tool_description
    assert "do not call save_assets" in tool_description


def test_canonical_tool_definitions_exclude_extract_assets_when_disabled() -> None:
    tool_names = [
        tool.name
        for tool in canonical_tool_definitions(True, asset_extraction_enabled=False)
    ]
    assert "extract_assets" not in tool_names


def test_provider_session_excludes_extract_assets_without_gemini_key() -> None:
    session = create_provider_session(
        model=Llm.GPT_5_5_HIGH,
        prompt_messages=[{"role": "user", "content": "Build a page."}],
        should_generate_images=True,
        openai_api_key="openai-key",
        openai_base_url=None,
        anthropic_api_key=None,
        gemini_api_key=None,
        replicate_api_key=None,
    )

    tools = cast(list[dict[str, Any]], getattr(session, "_tools"))
    tool_names = [tool["name"] for tool in tools]
    assert "extract_assets" not in tool_names


def test_provider_session_includes_extract_assets_with_gemini_key() -> None:
    session = create_provider_session(
        model=Llm.GPT_5_5_HIGH,
        prompt_messages=[{"role": "user", "content": "Build a page."}],
        should_generate_images=True,
        openai_api_key="openai-key",
        openai_base_url=None,
        anthropic_api_key=None,
        gemini_api_key="gemini-key",
        replicate_api_key=None,
    )

    tools = cast(list[dict[str, Any]], getattr(session, "_tools"))
    tool_names = [tool["name"] for tool in tools]
    assert "extract_assets" in tool_names


def test_provider_session_excludes_extract_assets_when_user_disables_it() -> None:
    session = create_provider_session(
        model=Llm.GPT_5_5_HIGH,
        prompt_messages=[{"role": "user", "content": "Build a page."}],
        should_generate_images=True,
        openai_api_key="openai-key",
        openai_base_url=None,
        anthropic_api_key=None,
        gemini_api_key="gemini-key",
        replicate_api_key=None,
        should_extract_assets=False,
    )

    tools = cast(list[dict[str, Any]], getattr(session, "_tools"))
    tool_names = [tool["name"] for tool in tools]
    assert "extract_assets" not in tool_names


def test_provider_session_excludes_screenshot_preview_when_chromium_unavailable(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(
        "agent.providers.factory.is_screenshot_preview_available", lambda: False
    )
    session = create_provider_session(
        model=Llm.GPT_5_5_HIGH,
        prompt_messages=[{"role": "user", "content": "Build a page."}],
        should_generate_images=True,
        openai_api_key="openai-key",
        openai_base_url=None,
        anthropic_api_key=None,
        gemini_api_key=None,
        replicate_api_key=None,
    )

    tools = cast(list[dict[str, Any]], getattr(session, "_tools"))
    tool_names = [tool["name"] for tool in tools]
    assert "screenshot_preview" not in tool_names


def test_provider_session_includes_screenshot_preview_when_chromium_available(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(
        "agent.providers.factory.is_screenshot_preview_available", lambda: True
    )
    session = create_provider_session(
        model=Llm.GPT_5_5_HIGH,
        prompt_messages=[{"role": "user", "content": "Build a page."}],
        should_generate_images=True,
        openai_api_key="openai-key",
        openai_base_url=None,
        anthropic_api_key=None,
        gemini_api_key=None,
        replicate_api_key=None,
    )

    tools = cast(list[dict[str, Any]], getattr(session, "_tools"))
    tool_names = [tool["name"] for tool in tools]
    assert "screenshot_preview" in tool_names


def test_extract_assets_tool_input_summary_uses_asset_descriptions() -> None:
    summary = summarize_tool_input(
        ToolCall(
            id="call-1",
            name="extract_assets",
            arguments={"asset_descriptions": ["logo", "avatar"]},
        ),
        AgentFileState(),
    )

    assert summary == {
        "count": 2,
        "asset_descriptions": ["logo", "avatar"],
    }


def test_canonical_tool_definitions_include_screenshot_preview() -> None:
    tools = canonical_tool_definitions(True)
    tool_names = [tool.name for tool in tools]
    assert "screenshot_preview" in tool_names
    screenshot_tool = next(
        tool for tool in tools if tool.name == "screenshot_preview"
    )
    assert screenshot_tool.parameters["properties"] == {}


def test_canonical_tool_definitions_exclude_screenshot_preview_when_disabled() -> None:
    tool_names = [
        tool.name
        for tool in canonical_tool_definitions(True, screenshot_enabled=False)
    ]
    assert "screenshot_preview" not in tool_names


def test_canonical_tool_definitions_include_batched_edit_images() -> None:
    tools = canonical_tool_definitions(True)
    tool_names = [tool.name for tool in tools]
    assert "edit_images" in tool_names
    edit_images_tool = next(tool for tool in tools if tool.name == "edit_images")
    assert "upscale" in edit_images_tool.description.lower()
    assert edit_images_tool.parameters["required"] == ["edits"]
    edits_schema = edit_images_tool.parameters["properties"]["edits"]
    assert edits_schema["type"] == "array"
    assert edits_schema["minItems"] == 1
    edit_schema = edits_schema["items"]
    assert edit_schema["required"] == ["prompt", "image_urls"]
    properties = edit_schema["properties"]
    assert properties["image_urls"]["type"] == "array"
    assert "turbo" not in properties
    assert "seed" not in properties
    assert properties["aspect_ratio"]["enum"] == list(P_IMAGE_EDIT_ASPECT_RATIOS)


def test_canonical_tool_definitions_exclude_edit_images_when_disabled() -> None:
    tool_names = [
        tool.name
        for tool in canonical_tool_definitions(True, image_editing_enabled=False)
    ]
    assert "edit_images" not in tool_names


def test_edit_images_tool_input_summary_preserves_edit_and_image_order() -> None:
    summary = summarize_tool_input(
        ToolCall(
            id="call-1",
            name="edit_images",
            arguments={
                "edits": [
                    {
                        "prompt": "Make image 1 monochrome",
                        "image_urls": [
                            "https://example.com/main.png",
                            "https://example.com/reference.png",
                        ],
                        "aspect_ratio": "16:9",
                    },
                    {
                        "prompt": "Upscale the logo",
                        "image_urls": ["https://example.com/logo.png"],
                    },
                ]
            },
        ),
        AgentFileState(),
    )

    assert summary == {
        "count": 2,
        "edits": [
            {
                "prompt": "Make image 1 monochrome",
                "image_urls": [
                    "https://example.com/main.png",
                    "https://example.com/reference.png",
                ],
                "aspect_ratio": "16:9",
            },
            {
                "prompt": "Upscale the logo",
                "image_urls": ["https://example.com/logo.png"],
                "aspect_ratio": "match_input_image",
            },
        ],
    }