vibe/tests/cli/textual_ui/test_image_input_validation.py

from __future__ import annotations

from pathlib import Path

import pytest

from tests.conftest import build_test_vibe_app, build_test_vibe_config
from vibe.cli.textual_ui.app import _ImageAttachmentRejection
from vibe.cli.textual_ui.widgets.chat_input.container import ChatInputContainer
from vibe.cli.textual_ui.widgets.messages import ErrorMessage
from vibe.core.autocompletion.path_prompt import build_path_prompt_payload
from vibe.core.config import ModelConfig, ProviderConfig
from vibe.core.types import (
    MAX_IMAGE_BYTES,
    MAX_IMAGES_PER_MESSAGE,
    Backend,
    ImageAttachment,
)

PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16


def _vision_config(*, supports_images: bool = True):
    models = [
        ModelConfig(
            name="mistral-vibe-cli-latest",
            provider="mistral",
            alias="devstral-latest",
            supports_images=supports_images,
        )
    ]
    providers = [
        ProviderConfig(
            name="mistral",
            api_base="https://api.mistral.ai/v1",
            api_key_env_var="MISTRAL_API_KEY",
            backend=Backend.MISTRAL,
        )
    ]
    return build_test_vibe_config(
        active_model="devstral-latest", models=models, providers=providers
    )


@pytest.mark.asyncio
async def test_build_image_attachments_happy_path(tmp_path: Path) -> None:
    (tmp_path / "shot.png").write_bytes(PNG_BYTES)
    payload = build_path_prompt_payload("look at @shot.png", base_dir=tmp_path)

    app = build_test_vibe_app(config=_vision_config())
    result = await app._build_image_attachments(payload)

    assert isinstance(result, list)
    assert len(result) == 1
    assert isinstance(result[0], ImageAttachment)
    assert result[0].alias == "shot.png"
    assert result[0].mime_type == "image/png"


@pytest.mark.asyncio
async def test_build_image_attachments_returns_empty_when_no_images(
    tmp_path: Path,
) -> None:
    (tmp_path / "notes.md").write_text("hi")
    payload = build_path_prompt_payload("read @notes.md", base_dir=tmp_path)

    app = build_test_vibe_app(config=_vision_config())
    result = await app._build_image_attachments(payload)

    assert result == []


@pytest.mark.asyncio
async def test_build_image_attachments_rejects_too_many_images(tmp_path: Path) -> None:
    mentions = []
    for i in range(MAX_IMAGES_PER_MESSAGE + 1):
        name = f"img{i}.png"
        (tmp_path / name).write_bytes(PNG_BYTES)
        mentions.append(f"@{name}")
    payload = build_path_prompt_payload(" ".join(mentions), base_dir=tmp_path)

    app = build_test_vibe_app(config=_vision_config())
    result = await app._build_image_attachments(payload)

    assert isinstance(result, _ImageAttachmentRejection)
    assert not result.no_vision
    assert "Too many image attachments" in result.message
    assert str(MAX_IMAGES_PER_MESSAGE) in result.message


@pytest.mark.asyncio
async def test_build_image_attachments_rejects_non_vision_model(tmp_path: Path) -> None:
    (tmp_path / "shot.png").write_bytes(PNG_BYTES)
    payload = build_path_prompt_payload("look at @shot.png", base_dir=tmp_path)

    app = build_test_vibe_app(config=_vision_config(supports_images=False))
    result = await app._build_image_attachments(payload)

    assert isinstance(result, _ImageAttachmentRejection)
    assert result.no_vision
    assert "does not support images" in result.message
    assert "devstral-latest" in result.message


@pytest.mark.asyncio
async def test_build_image_attachments_rejects_oversize_image(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    # Patch the cap down so we don't need to write 10MB to disk.
    monkeypatch.setattr("vibe.cli.textual_ui.app.MAX_IMAGE_BYTES", 32)

    (tmp_path / "shot.png").write_bytes(PNG_BYTES + b"\x00" * 64)
    payload = build_path_prompt_payload("look at @shot.png", base_dir=tmp_path)

    app = build_test_vibe_app(config=_vision_config())
    result = await app._build_image_attachments(payload)

    assert isinstance(result, _ImageAttachmentRejection)
    assert not result.no_vision
    assert "shot.png" in result.message
    assert "max" in result.message.lower()


def test_max_image_bytes_default_is_10_mib() -> None:
    assert MAX_IMAGE_BYTES == 10 * 1024 * 1024


def test_max_images_per_message_default_is_8() -> None:
    assert MAX_IMAGES_PER_MESSAGE == 8


@pytest.mark.asyncio
async def test_submit_restores_input_when_image_validation_fails(
    tmp_path: Path,
) -> None:
    (tmp_path / "shot.png").write_bytes(PNG_BYTES)
    app = build_test_vibe_app(config=_vision_config(supports_images=False))
    typed = f"look at @{tmp_path / 'shot.png'}"

    async with app.run_test() as pilot:
        chat_input = app.query_one(ChatInputContainer)
        chat_input.value = typed
        await pilot.press("enter")
        await pilot.pause()

        assert chat_input.value == typed
        errors = list(app.query(ErrorMessage))
        assert errors
        assert all(not e._show_border for e in errors)