Leonid Pershin
4186d0bcf1
- Introduced `fmt_bytes` function for human-readable byte size formatting. - Added `DownloadProgress` class to track and display download progress with speed and completion percentage. - Updated `download` functions in `civitai_fetch.py` and `llamacpp_fetch.py` to utilize progress tracking. - Enhanced `ollama_pull.py` to support streaming progress updates during model pulls. - Updated tests to validate new formatting and progress tracking functionalities, ensuring accurate output and user feedback.
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
from pathlib import Path
|
|
|
|
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
|