40 lines
881 B
Python
40 lines
881 B
Python
import pytest
|
|
|
|
from model_batch_downloader.progress import (
|
|
Aria2Progress,
|
|
format_bytes,
|
|
parse_aria2_progress,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("line", "expected"),
|
|
[
|
|
(
|
|
"[#abc 420MiB/1.0GiB(41%) CN:16 DL:18MiB ETA:33s]",
|
|
Aria2Progress(41, "18MiB", "33s"),
|
|
),
|
|
(
|
|
"[#abc 0B/0B(0%) CN:1 DL:0B]",
|
|
Aria2Progress(0, "0B", None),
|
|
),
|
|
("Download complete: /models/a.safetensors", None),
|
|
("", None),
|
|
],
|
|
)
|
|
def test_parse_aria2_progress(line, expected):
|
|
assert parse_aria2_progress(line) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[
|
|
(0, "0 B"),
|
|
(1024, "1.00 KiB"),
|
|
(1024**2 * 3, "3.00 MiB"),
|
|
(1024**3 * 2, "2.00 GiB"),
|
|
],
|
|
)
|
|
def test_format_bytes(value, expected):
|
|
assert format_bytes(value) == expected
|