zyt/app/tests/test_image_drop.py

from __future__ import annotations

import base64
import os
from concurrent.futures import ThreadPoolExecutor
from dataclasses import FrozenInstanceError
from pathlib import Path
from tempfile import TemporaryDirectory

os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")

import httpx
import pytest
import shiboken6
from PySide6.QtCore import QBuffer, QIODevice, QMimeData, QUrl
from PySide6.QtGui import QImage, QPixmap
from PySide6.QtWidgets import QApplication

from doctor_workstation.ui import image_drop
from doctor_workstation.ui.image_drop import (
    can_drop_images,
    capture_image_drop,
    materialize_image_drop,
)


@pytest.fixture(scope="module")
def application() -> QApplication:
    return QApplication.instance() or QApplication([])


def _image(width: int = 12, height: int = 8) -> QImage:
    image = QImage(width, height, QImage.Format.Format_RGB32)
    image.fill(0xFF668877)
    return image


def _image_bytes(image_format: str, width: int = 12, height: int = 8) -> bytes:
    buffer = QBuffer()
    assert buffer.open(QIODevice.OpenModeFlag.WriteOnly)
    assert _image(width, height).save(buffer, image_format)
    return bytes(buffer.data())


def _png(width: int = 12, height: int = 8) -> bytes:
    return _image_bytes("PNG", width, height)


def _local_mime(*paths: Path) -> QMimeData:
    mime = QMimeData()
    mime.setUrls([QUrl.fromLocalFile(str(path)) for path in paths])
    return mime


def _raw_mime(payload: bytes, media_type: str = "image/png") -> QMimeData:
    mime = QMimeData()
    mime.setData(media_type, payload)
    return mime


def _mock_download(monkeypatch: pytest.MonkeyPatch, handler: object) -> None:
    monkeypatch.setattr(
        image_drop,
        "_new_download_client",
        lambda: httpx.Client(transport=httpx.MockTransport(handler), follow_redirects=False),
    )


def test_local_images_are_validated_by_content_and_copied_without_touching_originals(
    tmp_path: Path,
) -> None:
    first = tmp_path / "第一张 图片.jpg"
    second = tmp_path / "second.png"
    first.write_bytes(_png())
    second.write_bytes(_png(10, 20))
    mime = _local_mime(first, second)
    assert can_drop_images(mime)
    drop = capture_image_drop(mime)
    assert drop.count == 2
    with materialize_image_drop(drop) as paths:
        assert len(paths) == 2
        assert all(path.suffix == ".png" for path in paths)
        assert paths[0].read_bytes() == first.read_bytes()
        assert paths[1].read_bytes() == second.read_bytes()
        assert all(path.parent != tmp_path for path in paths)
    assert all(not path.exists() for path in paths)
    assert first.exists() and second.exists()


@pytest.mark.parametrize("source_kind", ["local", "raw"])
def test_bmp_is_converted_to_server_supported_png(tmp_path: Path, source_kind: str) -> None:
    original = tmp_path / "image.bmp"
    content = _image_bytes("BMP")
    original.write_bytes(content)
    mime = _local_mime(original) if source_kind == "local" else _raw_mime(content, "image/bmp")
    with materialize_image_drop(capture_image_drop(mime)) as paths:
        assert paths[0].suffix == ".png"
        assert paths[0].read_bytes().startswith(b"\x89PNG\r\n\x1a\n")
        converted = QImage(str(paths[0]))
        assert converted.size() == _image().size()
        assert converted.pixelColor(0, 0) == _image().pixelColor(0, 0)
    assert original.read_bytes() == content
    assert not paths[0].exists()


def test_conversion_rechecks_byte_limit_and_cleans_up(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
) -> None:
    # A genuine 1-pixel BMP is smaller than its genuine PNG representation.
    content = _image_bytes("BMP", 1, 1)
    assert len(content) < len(_png(1, 1))
    monkeypatch.setattr(image_drop, "MAX_IMAGE_BYTES", len(content))
    monkeypatch.setattr(
        image_drop, "TemporaryDirectory",
        lambda prefix: TemporaryDirectory(prefix=prefix, dir=tmp_path),
    )
    drop = capture_image_drop(_raw_mime(content, "image/bmp"))
    with pytest.raises(ValueError, match="10 MiB"), materialize_image_drop(drop):
        pytest.fail("converted image exceeded upload limit")
    assert list(tmp_path.iterdir()) == []


def test_hover_never_reads_files_downloads_or_decodes(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
) -> None:
    def unexpected(*args: object, **kwargs: object) -> None:
        pytest.fail("hover must not access images")

    monkeypatch.setattr(Path, "is_file", unexpected)
    monkeypatch.setattr(image_drop, "_download_image", unexpected)
    monkeypatch.setattr(image_drop, "QImageReader", unexpected)
    assert can_drop_images(_local_mime(tmp_path / "not-yet-existing.png"))
    assert can_drop_images(_raw_mime(b"undecodable"))
    mime = QMimeData()
    mime.setHtml('<a href="https://site.invalid"><img src="https://img.invalid/a.png"></a>')
    assert can_drop_images(mime)


@pytest.mark.parametrize("pixmap", [False, True])
def test_memory_snapshot_survives_mime_deletion_and_works_in_worker(
    application: QApplication, pixmap: bool,
) -> None:
    original = _image()
    mime = QMimeData()
    mime.setImageData(QPixmap.fromImage(original) if pixmap else original)
    drop = capture_image_drop(mime)
    original.fill(0)
    shiboken6.delete(mime)
    assert drop.count == 1
    with pytest.raises(FrozenInstanceError):
        drop.sources = ()

    def resolve() -> bytes:
        with materialize_image_drop(drop) as paths:
            return paths[0].read_bytes()

    with ThreadPoolExecutor(max_workers=1) as worker:
        content = worker.submit(resolve).result(timeout=10)
    assert content == _png()


def test_local_files_take_precedence_over_duplicate_memory_picture(tmp_path: Path) -> None:
    path = tmp_path / "image.png"
    path.write_bytes(_png())
    mime = _local_mime(path)
    mime.setImageData(_image(3, 3))
    drop = capture_image_drop(mime)
    assert drop.count == 1
    with materialize_image_drop(drop) as paths:
        assert paths[0].read_bytes() == path.read_bytes()


def test_raw_image_snapshot_keeps_no_mime_reference() -> None:
    payload = _png()
    mime = _raw_mime(payload)
    drop = capture_image_drop(mime)
    mime.setData("image/png", b"modified")
    shiboken6.delete(mime)
    with materialize_image_drop(drop) as paths:
        assert paths[0].read_bytes() == payload


def test_browser_image_src_wins_over_enclosing_link_and_uri_list(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    urls: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        urls.append(str(request.url))
        assert "authorization" not in request.headers
        assert "cookie" not in request.headers
        return httpx.Response(200, content=_png(), headers={"Content-Type": "image/png"})

    _mock_download(monkeypatch, handler)
    mime = QMimeData()
    mime.setUrls([QUrl("https://site.invalid/article")])
    mime.setHtml(
        '<a href="https://site.invalid/article">'
        '<img src="https://cdn.invalid/a.png?width=2&amp;height=3"></a>'
    )
    with materialize_image_drop(capture_image_drop(mime)) as paths:
        assert len(paths) == 1
    assert urls == ["https://cdn.invalid/a.png?width=2&height=3"]


def test_multiple_html_images_and_base64_data_are_preserved() -> None:
    payload = _png()
    encoded = base64.b64encode(payload).decode("ascii")
    mime = QMimeData()
    mime.setHtml(f'<img src="data:image/png;base64,{encoded}">' * 2)
    drop = capture_image_drop(mime)
    assert drop.count == 2
    with materialize_image_drop(drop) as paths:
        assert [path.read_bytes() for path in paths] == [payload, payload]


@pytest.mark.parametrize("representation", ["text", "url"])
def test_remote_image_url_can_be_dropped(
    monkeypatch: pytest.MonkeyPatch, representation: str,
) -> None:
    _mock_download(monkeypatch, lambda _: httpx.Response(200, content=_png()))
    mime = QMimeData()
    if representation == "url":
        mime.setUrls([QUrl("https://cdn.invalid/image?id=123")])
    else:
        mime.setText("https://cdn.invalid/image?id=123")
    assert can_drop_images(mime)
    with materialize_image_drop(capture_image_drop(mime)) as paths:
        assert paths[0].read_bytes() == _png()


@pytest.mark.parametrize("text", ["C:/private/photo.png", "file:///C:/private/photo.png", "hello"])
def test_plain_text_is_not_interpreted_as_local_file(text: str) -> None:
    mime = QMimeData()
    mime.setText(text)
    assert not can_drop_images(mime)
    with pytest.raises(ValueError, match="请拖入图片"):
        capture_image_drop(mime)


def test_mixed_valid_and_invalid_files_fail_before_yield_and_clean_up(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
) -> None:
    good, bad = tmp_path / "good.png", tmp_path / "fake.png"
    good.write_bytes(_png())
    bad.write_text("not a real image", encoding="utf-8")
    monkeypatch.setattr(
        image_drop, "TemporaryDirectory",
        lambda prefix: TemporaryDirectory(prefix=prefix, dir=tmp_path),
    )
    with (
        pytest.raises(ValueError, match="不是受支持的图片"),
        materialize_image_drop(capture_image_drop(_local_mime(good, bad))),
    ):
        pytest.fail("must validate the full batch before uploading anything")
    assert sorted(path.name for path in tmp_path.iterdir()) == ["fake.png", "good.png"]


def test_temp_files_are_cleaned_when_upload_raises() -> None:
    with (
        pytest.raises(RuntimeError, match="upload failed"),
        materialize_image_drop(capture_image_drop(_raw_mime(_png()))) as paths,
    ):
        raise RuntimeError("upload failed")
    assert not paths[0].parent.exists()


@pytest.mark.parametrize("kind", ["local", "html"])
def test_batch_limit_is_checked_without_io(tmp_path: Path, kind: str) -> None:
    if kind == "local":
        mime = _local_mime(*(tmp_path / f"{index}.png" for index in range(21)))
    else:
        mime = QMimeData()
        mime.setHtml('<img src="https://cdn.invalid/a.png">' * 21)
    with pytest.raises(ValueError, match="20"):
        capture_image_drop(mime)


def test_exactly_twenty_images_are_allowed(tmp_path: Path) -> None:
    path = tmp_path / "image.png"
    path.write_bytes(_png())
    with materialize_image_drop(capture_image_drop(_local_mime(*([path] * 20)))) as paths:
        assert len(paths) == 20


def test_local_byte_limit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(image_drop, "MAX_IMAGE_BYTES", 100)
    path = tmp_path / "large.png"
    path.write_bytes(b"a" * 101)
    with (
        pytest.raises(ValueError, match="10 MiB"),
        materialize_image_drop(capture_image_drop(_local_mime(path))),
    ):
        pytest.fail("oversized image accepted")


def test_raw_byte_limit_is_checked_during_capture(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(image_drop, "MAX_IMAGE_BYTES", 10)
    with pytest.raises(ValueError, match="10 MiB"):
        capture_image_drop(_raw_mime(b"a" * 11))


def test_pixel_limit_is_checked_for_encoded_and_memory_images(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(image_drop, "MAX_IMAGE_PIXELS", 50)
    with (
        pytest.raises(ValueError, match="像素过大"),
        materialize_image_drop(capture_image_drop(_raw_mime(_png()))),
    ):
        pytest.fail("large raster accepted")
    mime = QMimeData()
    mime.setImageData(_image())
    with pytest.raises(ValueError, match="像素过大"):
        capture_image_drop(mime)


@pytest.mark.parametrize("payload", [
    b'<svg xmlns="http://www.w3.org/2000/svg" width="10" height="10"></svg>',
    b"<html><body>login</body></html>",
    b"not an image",
])
def test_non_raster_content_cannot_masquerade_as_png(payload: bytes) -> None:
    with (
        pytest.raises(ValueError, match="不是受支持的图片"),
        materialize_image_drop(capture_image_drop(_raw_mime(payload))),
    ):
        pytest.fail("invalid image accepted")


@pytest.mark.parametrize("value", [
    "data:image/svg+xml;base64,PHN2Zy8+",
    "data:image/png;base64,not-base64!",
    "data:image/png,not-base64",
])
def test_invalid_data_images_fail(value: str) -> None:
    mime = QMimeData()
    mime.setText(value)
    with pytest.raises(ValueError), materialize_image_drop(capture_image_drop(mime)):
        pytest.fail("invalid data image accepted")


@pytest.mark.parametrize("redirect", ["file:///C:/private/image.png", "javascript:alert(1)"])
def test_unsafe_redirects_are_never_requested(
    monkeypatch: pytest.MonkeyPatch, redirect: str,
) -> None:
    requests: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(str(request.url))
        return httpx.Response(302, headers={"Location": redirect})

    _mock_download(monkeypatch, handler)
    mime = QMimeData()
    mime.setText("https://cdn.invalid/a.png")
    with (
        pytest.raises(ValueError, match="HTTP/HTTPS"),
        materialize_image_drop(capture_image_drop(mime)),
    ):
        pytest.fail("unsafe redirect accepted")
    assert requests == ["https://cdn.invalid/a.png"]


def test_redirects_do_not_forward_cookies(monkeypatch: pytest.MonkeyPatch) -> None:
    requests: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(str(request.url))
        assert "cookie" not in request.headers
        assert "authorization" not in request.headers
        if len(requests) == 1:
            return httpx.Response(302, headers={"Location": "/b.png", "Set-Cookie": "token=secret"})
        return httpx.Response(200, content=_png())

    _mock_download(monkeypatch, handler)
    mime = QMimeData()
    mime.setText("https://cdn.invalid/a.png")
    with materialize_image_drop(capture_image_drop(mime)) as paths:
        assert paths[0].read_bytes() == _png()
    assert len(requests) == 2


@pytest.mark.parametrize("status,headers", [
    (401, {}),
    (403, {}),
    (200, {"Content-Type": "text/html"}),
])
def test_login_and_webpage_sources_explain_save_locally(
    monkeypatch: pytest.MonkeyPatch, status: int, headers: dict[str, str],
) -> None:
    _mock_download(monkeypatch, lambda _: httpx.Response(status, headers=headers, content=b"login"))
    mime = QMimeData()
    mime.setText("https://site.invalid/private.png")
    with (
        pytest.raises(ValueError, match="先保存本地再拖入"),
        materialize_image_drop(capture_image_drop(mime)),
    ):
        pytest.fail("login page accepted")


@pytest.mark.parametrize("declared_length", [True, False])
def test_download_size_limit_uses_headers_and_actual_stream(
    monkeypatch: pytest.MonkeyPatch, declared_length: bool,
) -> None:
    monkeypatch.setattr(image_drop, "MAX_IMAGE_BYTES", 100)

    class OversizedStream(httpx.SyncByteStream):
        def __iter__(self):
            yield b"a" * 60
            yield b"b" * 60

    def handler(_: httpx.Request) -> httpx.Response:
        headers = {"Content-Length": "120"} if declared_length else {}
        return httpx.Response(200, headers=headers, stream=OversizedStream())

    _mock_download(monkeypatch, handler)
    mime = QMimeData()
    mime.setText("https://cdn.invalid/a.png")
    with (
        pytest.raises(ValueError, match="10 MiB"),
        materialize_image_drop(capture_image_drop(mime)),
    ):
        pytest.fail("oversized response accepted")


def test_embedded_url_credentials_are_rejected() -> None:
    mime = QMimeData()
    mime.setText("https://token:secret@cdn.invalid/a.png")
    with pytest.raises(ValueError, match="HTTP/HTTPS"):
        capture_image_drop(mime)