62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import io
|
|
import pytest
|
|
import torch
|
|
from unittest.mock import MagicMock
|
|
from PIL import Image
|
|
import pixiv_node
|
|
from pixiv_node import PixivBrowser
|
|
|
|
|
|
def _make_mock_client(w=64, h=64):
|
|
client = MagicMock()
|
|
img = Image.new("RGB", (w, h), color=(128, 64, 200))
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
client.download_image_bytes.return_value = buf.getvalue()
|
|
client.get_original_url.return_value = "https://i.pximg.net/orig/img.jpg"
|
|
return client
|
|
|
|
|
|
def test_execute_returns_image_tensor(monkeypatch):
|
|
monkeypatch.setattr(pixiv_node, "_get_client", lambda: _make_mock_client())
|
|
result = PixivBrowser().execute(artwork_ids="12345")
|
|
tensor = result[0]
|
|
assert isinstance(tensor, torch.Tensor)
|
|
assert tensor.ndim == 4
|
|
assert tensor.shape[0] == 1
|
|
assert tensor.shape[3] == 3
|
|
assert tensor.dtype == torch.float32
|
|
assert tensor.min() >= 0.0 and tensor.max() <= 1.0
|
|
|
|
|
|
def test_execute_multiple_ids_returns_batch(monkeypatch):
|
|
monkeypatch.setattr(pixiv_node, "_get_client", lambda: _make_mock_client())
|
|
result = PixivBrowser().execute(artwork_ids="111,222,333")
|
|
assert result[0].shape[0] == 3
|
|
|
|
|
|
def test_execute_skips_failed_downloads(monkeypatch):
|
|
client = MagicMock()
|
|
client.get_original_url.return_value = "https://i.pximg.net/img.jpg"
|
|
img = Image.new("RGB", (32, 32))
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
client.download_image_bytes.side_effect = [Exception("timeout"), buf.getvalue()]
|
|
monkeypatch.setattr(pixiv_node, "_get_client", lambda: client)
|
|
result = PixivBrowser().execute(artwork_ids="bad_id,good_id")
|
|
assert result[0].shape[0] == 1
|
|
|
|
|
|
def test_execute_raises_on_empty_ids(monkeypatch):
|
|
monkeypatch.setattr(pixiv_node, "_get_client", lambda: MagicMock())
|
|
with pytest.raises(ValueError, match="请先在弹窗中选择图片"):
|
|
PixivBrowser().execute(artwork_ids="")
|
|
|
|
|
|
def test_input_types_has_artwork_ids():
|
|
assert "artwork_ids" in PixivBrowser.INPUT_TYPES()["required"]
|
|
|
|
|
|
def test_return_types_is_image():
|
|
assert PixivBrowser.RETURN_TYPES == ("IMAGE",)
|