141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
import json
|
|
|
|
import pytest
|
|
|
|
import model_batch_downloader.nodes as subject
|
|
|
|
|
|
def sample_manifest():
|
|
return json.dumps(
|
|
[
|
|
{
|
|
"url": "https://example.com/a.safetensors",
|
|
"model_type": "checkpoints",
|
|
"filename": "a.safetensors",
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
def test_roots_preserve_registered_order_and_use_private_defaults(
|
|
monkeypatch, tmp_path
|
|
):
|
|
models_dir = tmp_path / "models"
|
|
first = tmp_path / "first-controlnet"
|
|
second = tmp_path / "second-controlnet"
|
|
registered = {
|
|
"controlnet": ([str(first), str(second)], {".safetensors"}),
|
|
}
|
|
|
|
monkeypatch.setattr(subject.folder_paths, "models_dir", str(models_dir))
|
|
monkeypatch.setattr(subject.folder_paths, "folder_names_and_paths", registered)
|
|
|
|
def get_folder_paths(name):
|
|
paths, _extensions = registered[name]
|
|
return paths
|
|
|
|
monkeypatch.setattr(subject.folder_paths, "get_folder_paths", get_folder_paths)
|
|
|
|
roots = subject._roots()
|
|
|
|
assert roots["controlnet"] == (str(first), str(second))
|
|
assert roots["onnx"] == (str(models_dir / "onnx"),)
|
|
assert roots["sam3"] == (str(models_dir / "sam3"),)
|
|
assert roots["llm"] == (str(models_dir / "llm"),)
|
|
assert roots["ultralytics_bbox"] == (
|
|
str(models_dir / "ultralytics" / "bbox"),
|
|
)
|
|
assert roots["ultralytics_segm"] == (
|
|
str(models_dir / "ultralytics" / "segm"),
|
|
)
|
|
|
|
|
|
def test_both_nodes_share_executor_and_preserve_identity(monkeypatch):
|
|
sentinel = object()
|
|
result = object()
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
subject,
|
|
"execute_manifest",
|
|
lambda manifest_json, passthrough=None: calls.append(manifest_json)
|
|
or (passthrough, result),
|
|
)
|
|
|
|
assert subject.ModelBatchDownloader().download(sample_manifest(), sentinel) == (
|
|
sentinel,
|
|
result,
|
|
)
|
|
assert subject.ModelBatchDownloaderJSON().download(sample_manifest(), sentinel) == (
|
|
sentinel,
|
|
result,
|
|
)
|
|
assert calls == [sample_manifest(), sample_manifest()]
|
|
|
|
|
|
def test_authoring_nodes_are_output_nodes_and_always_change():
|
|
for node_class in (
|
|
subject.ModelBatchDownloader,
|
|
subject.ModelBatchDownloaderJSON,
|
|
):
|
|
assert node_class.OUTPUT_NODE is True
|
|
assert node_class.RETURN_NAMES == ("passthrough", "download_result")
|
|
assert str(node_class.RETURN_TYPES[1]) == "DOWNLOAD_RESULT"
|
|
assert node_class.IS_CHANGED(sample_manifest()) != node_class.IS_CHANGED(
|
|
sample_manifest()
|
|
)
|
|
|
|
|
|
def test_table_and_json_nodes_expose_same_backend_schema():
|
|
assert subject.ModelBatchDownloader.INPUT_TYPES() == (
|
|
subject.ModelBatchDownloaderJSON.INPUT_TYPES()
|
|
)
|
|
schema = subject.ModelBatchDownloader.INPUT_TYPES()
|
|
assert schema["required"]["manifest_json"][0] == "STRING"
|
|
assert "passthrough" in schema["optional"]
|
|
|
|
|
|
def test_execute_manifest_fails_before_parsing_when_aria2_missing(monkeypatch):
|
|
monkeypatch.setattr(subject.shutil, "which", lambda _name: None)
|
|
with pytest.raises(RuntimeError, match="aria2c"):
|
|
subject.execute_manifest("[]")
|
|
|
|
|
|
def test_execute_manifest_connects_runner_to_comfy_progress(monkeypatch):
|
|
progress = []
|
|
interruption_checks = []
|
|
resolved = (object(), object())
|
|
download_result = object()
|
|
|
|
class ProgressBar:
|
|
def __init__(self, total):
|
|
assert total == 200
|
|
|
|
def update_absolute(self, value, total=None):
|
|
progress.append((value, total))
|
|
|
|
monkeypatch.setattr(subject.comfy.utils, "ProgressBar", ProgressBar)
|
|
monkeypatch.setattr(subject, "parse_manifest", lambda _json: (object(),))
|
|
monkeypatch.setattr(subject, "resolve_manifest", lambda *_args: resolved)
|
|
monkeypatch.setattr(subject.shutil, "which", lambda _name: "aria2c")
|
|
monkeypatch.setattr(
|
|
subject.comfy.model_management,
|
|
"throw_exception_if_processing_interrupted",
|
|
lambda: interruption_checks.append(True),
|
|
)
|
|
|
|
def fake_run(items, executable, environ, progress_callback, check_interrupted):
|
|
assert items is resolved
|
|
assert executable == "aria2c"
|
|
progress_callback(75, 200)
|
|
check_interrupted()
|
|
return download_result
|
|
|
|
monkeypatch.setattr(subject, "run_downloads", fake_run)
|
|
|
|
assert subject.execute_manifest(sample_manifest(), roots_by_type={}) == (
|
|
None,
|
|
download_result,
|
|
)
|
|
assert progress == [(75, 200)]
|
|
assert interruption_checks == [True]
|