Leonid Pershin
506369548c
- Introduced the `pick_preview_image` function to extract the first usable preview image URL and its suffix from model version data. - Updated the `seed_civitai` function to include preview image URLs and destinations in job definitions, improving model handling. - Implemented the `ensure_preview` function to download missing preview images during job processing, enhancing user experience. - Added tests for `pick_preview_image` to ensure correct functionality across various scenarios, ensuring robustness in image handling.
83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
from pathlib import Path
|
|
|
|
from gpu_rent.civitai import pick_preview_image
|
|
from gpu_rent.remote.civitai_fetch import fmt_bytes, progress_line, should_skip
|
|
|
|
|
|
def test_should_skip_missing(tmp_path: Path):
|
|
skip, _ = should_skip(tmp_path / "nope.safetensors", "abc")
|
|
assert skip is False
|
|
|
|
|
|
def test_should_skip_exists_no_sha(tmp_path: Path):
|
|
dest = tmp_path / "a.safetensors"
|
|
dest.write_bytes(b"weights")
|
|
skip, reason = should_skip(dest, "")
|
|
assert skip is True
|
|
assert "уже есть" in reason
|
|
|
|
|
|
def test_should_skip_sha_match(tmp_path: Path):
|
|
dest = tmp_path / "a.safetensors"
|
|
dest.write_bytes(b"weights")
|
|
import hashlib
|
|
|
|
digest = hashlib.sha256(b"weights").hexdigest()
|
|
skip, reason = should_skip(dest, digest)
|
|
assert skip is True
|
|
assert "sha ok" in reason
|
|
|
|
|
|
def test_should_not_skip_bad_sha(tmp_path: Path):
|
|
dest = tmp_path / "a.safetensors"
|
|
dest.write_bytes(b"weights")
|
|
skip, reason = should_skip(dest, "deadbeef" * 8)
|
|
assert skip is False
|
|
assert "перекачиваю" in reason
|
|
|
|
|
|
def test_fmt_bytes():
|
|
assert fmt_bytes(500) == "500B"
|
|
assert "KB" in fmt_bytes(2048)
|
|
assert "MB" in fmt_bytes(5 * 1024 * 1024)
|
|
assert "GB" in fmt_bytes(2 * 1024**3)
|
|
|
|
|
|
def test_progress_line_with_total():
|
|
line = progress_line("file", 512 * 1024 * 1024, 1024 * 1024 * 1024, 10 * 1024 * 1024, width=10)
|
|
assert "[" in line and "]" in line
|
|
assert "50.0%" in line
|
|
assert "/s" in line
|
|
assert "file" in line
|
|
|
|
|
|
def test_progress_line_unknown_total():
|
|
line = progress_line("file", 1024 * 1024, None, 100_000)
|
|
assert "%" not in line
|
|
assert "/s" in line
|
|
|
|
|
|
def test_pick_preview_image_first_jpeg():
|
|
version = {
|
|
"images": [
|
|
{"type": "video", "url": "https://example.com/x.mp4"},
|
|
{"type": "image", "url": "https://image.civitai.com/x/width=450/abc.jpeg"},
|
|
]
|
|
}
|
|
got = pick_preview_image(version)
|
|
assert got is not None
|
|
url, suffix = got
|
|
assert "abc.jpeg" in url
|
|
assert suffix == ".preview.jpg"
|
|
|
|
|
|
def test_pick_preview_image_png_and_default():
|
|
assert pick_preview_image({"images": [{"url": "https://cdn.example/a.png"}]})[1] == (
|
|
".preview.png"
|
|
)
|
|
assert pick_preview_image({"images": [{"url": "https://cdn.example/hash/width=450"}]})[
|
|
1
|
|
] == ".preview.jpg"
|
|
assert pick_preview_image({"images": []}) is None
|
|
assert pick_preview_image({}) is None
|